For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Python function
add_layer_hook
add_layer_hook()
max.nn.layer.add_layer_hook(fn)
Adds a hook to call a function after each layer’s __call__.
The function will be passed four inputs:
- layer
- input_args
- input_kwargs
- outputs
The function can either return None or new outputs that will replace the layer returned outputs.
Note that input and outputs contain graph Values, which show limited
information (like shape and dtype). You can still see the computed values
if you include the Value in the graph.ops.output op, or call graph.ops.print.
Example of printing debug inputs:
from max.nn.layer import add_layer_hook
def print_info(layer, args, kwargs, outputs):
print("Layer:", type(layer).__name__)
print("Input args:", args)
print("Input kwargs:", kwargs)
print("Outputs:", outputs)
return outputs
add_layer_hook(print_info)