xlnstorch.viz.graph.make_autograd_graph

xlnstorch.viz.graph.make_autograd_graph#

xlnstorch.viz.graph.make_autograd_graph(*vars, graph_name='Autograd Graph', show_saved=False, leaf_color='orange', node_color='lightgrey', output_color='yellow', params=None, node_attr=None, edge_attr=None)#

Build (and return) a graphviz.Digraph object that visualizes the PyTorch autograd graph for the given variables.

Parameters:

  • vars (torch.Tensor | LNSTensor) – The output variables for which to build the autograd graph. Typically just the one output of the loss function.

  • graph_name (str, optional) – The name of the graph to be displayed in the visualization. Defaults to “Autograd Graph”.

  • show_saved (bool, optional) – If True, saved tensors and variables will be shown in the graph (inside their respective function nodes).

  • leaf_color (str, optional) – The color to use for leaf nodes (i.e., tensors that are not outputs of any autograd function). Defaults to “orange”.

  • node_color (str, optional) – The color to use for function nodes in the graph. Defaults to “lightgrey”.

  • output_color (str, optional) – The color to use for the final output tensor node in the graph. Defaults to “yellow”.

  • params (Dict[str, torch.Tensor | LNSTensor], optional) – An optional mapping parameter_name -> value where value is either a torch.Tensor or an LNSTensor instance. If supplied, the corresponding nodes will be highlighted and annotated with the user-provided name, which makes it much easier to see where model parameters occur in the graph.

  • node_attr (Dict[str, str], optional) – Additional attributes to apply to all nodes in the graph. This can be used to set styles, colors, or other properties that should be consistent across all nodes.

  • edge_attr (Dict[str, str], optional) – Additional attributes to apply to all edges in the graph. This can be used to set styles, colors, or other properties that should be consistent across all edges.

Returns:

A directed graph object representing the autograd graph. This can be rendered using graphviz.render or similar methods.

Return type:

graphviz.Digraph

Contents