View Source Nx.Defn.Tree (Nx v0.2.1)

Helper functions to traverse expressions, either as single nodes or recursively.

Link to this section Summary

Functions

Applies the given function to the arguments of the node, with the given accumulator as a starting value.

Check if the given tree has any of the given hooks in it.

Puts new args in the given tensor expression and gives it a new id.

Replaces args in the given tensor expression.

Rewrites the types of the given tensor expressions according to the given options.

Link to this section Functions

Link to this function

apply_args(expr, acc, fun)

View Source

Applies the given function to the arguments of the node, with the given accumulator as a starting value.

Warning: be very careful when using this function to traverse the expression recursively. If you plan to do so, you should consider also storing the visited nodes to avoid multiple traversals.

Check if the given tree has any of the given hooks in it.

Puts new args in the given tensor expression and gives it a new id.

Replaces args in the given tensor expression.

Use this function with extreme care. Changing the args but keeping the same id may mean you have different versions of the same node. Do this change only if you guarante all nodes in the tree have been replaced equally.

Link to this function

rewrite_types(tensor_expr, opts \\ [])

View Source

Rewrites the types of the given tensor expressions according to the given options.

options

Options

  • :max_float_type - set the max float type
  • :max_signed_type - set the max signed integer type
  • :max_unsigned_type - set the max unsigned integer type