Module graph 

The graph-level API: primitives, fused and custom nodes, and access to an array’s node.

This is what graph transformations such as nuvec-compile build on; everyday code never needs it. It is not covered by the API’s stability pins, and its functions return Result (they are not array operations: recording a primitive validates it, see from_primitive).

Re-exports

pub use crate::fast::FastOp;

Structs

CustomCall

One call of a custom function, as recorded in the graph: its rules and the number of primal inputs. Equal (and hashed) by the identity of the function.

Flags

How an array’s elements are laid out in memory.

Fused

A graph of element-wise operations evaluated as one kernel, with one or more of its steps as results. Cheap to clone; equal graphs compare and hash equal.

SliceArg

One slice specification per dimension: start:stop:step with Python semantics.

Enums

ArgReduceOp

The position of an extreme value along one axis.

BinaryOp

An element-wise function of two operands.

FusedRef

An operand of a fused step: an input of the fused computation or an earlier step’s result.

Primitive

An operation recorded in the graph, with its parameters.

ReduceOp

A reduction over a set of axes.

ScanOp

The combination a cumulative scan applies.

ScatterReduce

How a scatter combines an update with the value it lands on.

UnaryOp

An element-wise function of one operand.

Functions

broadcast_shapes

The broadcast shape of a and b (NumPy rules), or op’s shape-mismatch error, as the binary operations compute it.

flags

Memory-layout flags of a.

from_primitive

Records primitive applied to inputs, with the same validation as the operations, and returns one array per output. Failure is sticky: if an input has failed, the outputs fail with its error.

id

Identifies the operation that produces a (shared by its siblings) for as long as the array is alive.

inputs

The arrays a is computed from, until it has been evaluated. Empty for host data and evaluated arrays.

is_tracer

Whether a depends on a transform’s tracer or a placeholder, that is, it is being traced by grad, vjp, jvp, vmap or compile. Tracer status is transient: an array is a tracer only while the outermost transform it depends on runs (on this thread).

output_index

Which output of its operation a is.

placeholder

A placeholder with this dtype and shape, standing for an input while a function is traced (as in MLX): operations can be recorded on it, but evaluating anything that depends on it fails. Placeholders are tracers.

primitive

The operation that produces a, until it has been evaluated. None for host data and evaluated arrays.

siblings

Every output of the operation that produces a, in order (a is the one at output_index).

slice

a[args] with one SliceArg per dimension (Python semantics; a view), the lowered form of basic indexing.

strides

Per-dimension steps of a’s layout, in elements.