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

Helper functions to traverse defn expressions, either as single nodes or in-depth.

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.

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

Gets all IDs of all elements in the same scope.

Functions

Link to this function

apply_args(expr, type \\ :all, acc, fun)

View Source

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

By default, type is :all, which means all arguments are traversed. If type is :scope, only expressions that are in the same scope are traversed. Therefore, expressions such as while's condition and body, optional's default implementation, functions, and so forth are not traversed. Note conds are always traversed because, while they introduce a new scope, they can also access its parents directly, so you must take conds into account accordingly.

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 by using tensor.data.expr.id as cache key.

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.

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

  • :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
Link to this function

scope_ids(expr, ids \\ %{})

View Source

Gets all IDs of all elements in the same scope.

while's condition and body, fun's body and similar are considered different scopes. When it comes to cond, an ID will only be considered if it is used outside of the cond or used in several distinct conds. Constants are also ignored, as they have global IDs based on the constants themselves.

An existing maps of ids can be given to accumulate on top of it.