View Source Nx.Defn.Expr (Nx v0.6.2)

Summary

Functions

The expression used by Nx.Defn.Compiler.

Creates a cond tensor expression.

Creates a tensor expression metadata node wrapping the given tensor expression.

Creates a tensor expression parameter at pos based on the given tensor expression.

Creates a tensor expression parameter at pos based on the given tensor and context.

Creates a tensor expression parameter at pos with the given context, type, shape, and pos.

Builds an tensor expression from the given tensor.

Creates a tuple with elements in list that points to tuple expression expr.

Creates a while tensor expression.

Functions

Link to this function

%Nx.Defn.Expr{}

View Source (struct)

The expression used by Nx.Defn.Compiler.

Nx.Defn.Compiler changes Nx default backend from Nx.BinaryBackend to Nx.Defn.Expr. It is a struct with the following fields:

  • :id - a unique identifier
  • :op - the operation name
  • :args - the operation arguments
  • :context - the context of the expression. The default context is :root.

Convenience functions for traversing expressions and composite types can be found in Nx.Defn.Composite and Nx.Defn.Tree.

Syntax nodes

Most nodes are created directly via the Nx module and therefore map directly to Nx.Tensor callbacks. However the following syntax nodes exist:

  • parameter(integer)

  • constant(number)

  • tensor(tensor)

  • metadata(expr, metadata)

  • elem(tuple, pos) - created automatically from expression that return tuples. Note it may return tuples too, which means we have nested tuples

  • fun(parameters, t, mfa) - the mfa is used only for introspection purposes

  • cond(clauses, otherwise)

  • while(initial, condition, body)

  • attach_token(token(%Nx.Defn.Token{}), expr)

defn compilers must handle said nodes accordingly.

Creates a cond tensor expression.

Link to this function

metadata(expr, metadata)

View Source

Creates a tensor expression metadata node wrapping the given tensor expression.

The metadata is map. If the inspect key is present, it will be used to annotate the metadata when inspected. Otherwise the metadata node does not appear during inspection.

Creates a tensor expression parameter at pos based on the given tensor expression.

Link to this function

parameter(tensor, context, pos)

View Source

Creates a tensor expression parameter at pos based on the given tensor and context.

Link to this function

parameter(context, type, shape, pos)

View Source

Creates a tensor expression parameter at pos with the given context, type, shape, and pos.

Builds an tensor expression from the given tensor.

Creates a tuple with elements in list that points to tuple expression expr.

list must be a list of tensor expressions of the same size as the tuple expression.

Link to this function

while(initial, context, arg, condition, body)

View Source

Creates a while tensor expression.