ONNX operators as Nx calls. Values known at trace time live on Nx.BinaryBackend ("static"), so shape arithmetic folds to plain integers while tracing. Nx has no zero-size tensors, so a zero-size value is {:empty, shape}.
Summary
Functions
Computes the Softmax, LogSoftmax or Hardmax of node along its axis.
Returns the shape that ONNX broadcasting gives for shapes.
Returns whether EMLX computes tensor: in a function that EMLX compiles, or eagerly on EMLX.Backend.
Normalizes x over axes with fun, a kernel of EMLX.Fast that normalizes the last axis.
Returns the host value of value, or raises Transmuter.UnknownValueError when the value is not known on the host. A literal of a traced function and a large value that it folds from static values give their host tensor, and any other known value gives itself, so callers read the value that this function returns.
Returns whether the host can read a value.
Returns how EMLX.Fast.scaled_dot_product_attention reaches a fused kernel of MLX 0.32.0 for attention of these shapes: q {batch, heads, s, d}, k {batch, kv_heads, t, d} and a v head size dv. mask is :none, :array or :causal.
Computes (x - mean) / sqrt(variance + eps) over axes.
Reshapes t to rank rank, with its values along axis and size 1 on the other axes.
Reads the data of a tensor that keeps it in a file next to the model.
Rounds half to even, as ONNX does, on every backend.
Returns the shape of a tensor or of a zero-size value.
Splits t along axis into parts of sizes. A part of size 0 is a zero-size value.
Returns whether a value is known before the graph runs.
The shape and the type of the tensor that tensor/2 gives for t, from its dims and data_type, without its data. A string or an empty tensor gives nil.
Calls fun with values, where f16 and bf16 tensors become f32, and casts the f32 tensors that fun returns back to that type. fun takes and returns a list.
Functions
Computes the Softmax, LogSoftmax or Hardmax of node along its axis.
These operators work on one axis from opset 13. Before opset 13, they flatten the input to 2-D at the axis and work on the second dimension.
Returns the shape that ONNX broadcasting gives for shapes.
A dimension of 0 takes the place of 1 like any other size.
Returns whether EMLX computes tensor: in a function that EMLX compiles, or eagerly on EMLX.Backend.
Operators then use the fused kernels of EMLX.Fast. Other compilers, such as Nx.Defn.Evaluator, keep the plain Nx operations, also on EMLX tensors.
Normalizes x over axes with fun, a kernel of EMLX.Fast that normalizes the last axis.
The normalized axes become one axis, and each weight becomes a vector of this axis. fun gets this view of x and the weights.
Returns the host value of value, or raises Transmuter.UnknownValueError when the value is not known on the host. A literal of a traced function and a large value that it folds from static values give their host tensor, and any other known value gives itself, so callers read the value that this function returns.
Every read of a tensor value on the host goes through here, so a traced value raises the one error that Transmuter.compile/3 breaks the graph on.
Returns whether the host can read a value.
Concrete values can be read on the host: static ones, literals and constants of a traced function, large values that it folds from static values, and every tensor when the model runs eagerly. Other values inside a traced function are unknown.
Returns how EMLX.Fast.scaled_dot_product_attention reaches a fused kernel of MLX 0.32.0 for attention of these shapes: q {batch, heads, s, d}, k {batch, kv_heads, t, d} and a v head size dv. mask is :none, :array or :causal.
MLX decides in ScaledDotProductAttention::use_fallback (mlx/backend/metal/scaled_dot_product_attention.cpp:591-641 of MLX 0.32.0). The "vector" kernel takes s <= 8, s <= t, s * heads / kv_heads <= 32, and a head size of 64, 96, 128 or 256 (or d = 192 with dv = 128). The "full" kernel takes s > 8, d = dv of 64, 80 or 128, and no mask, an array mask, or a causal mask with s <= t. Other shapes take the fallback of mlx/fast.cpp, which builds the scores of all queries and keys at once.
Returns :fused when a kernel takes the shapes, {:pad, size} when q, k and v with zeros up to the head size size take a kernel and size is at most 1.25 times d, or :plain otherwise. Zeros after the head size do not change q * k, and the values of the zero columns of v are sliced off.
Computes (x - mean) / sqrt(variance + eps) over axes.
It centers x once and reuses it, where Nx.variance would compute the mean and center x again.
Reshapes t to rank rank, with its values along axis and size 1 on the other axes.
Reads the data of a tensor that keeps it in a file next to the model.
Rounds half to even, as ONNX does, on every backend.
Backends round ties differently: MLX to even, Nx.BinaryBackend away from zero. A tie becomes twice its rounded half, which is never a tie, so the rounding mode of the backend does not matter. Infinities and NaN pass through, because Nx.BinaryBackend cannot round them.
Returns the shape of a tensor or of a zero-size value.
Splits t along axis into parts of sizes. A part of size 0 is a zero-size value.
Returns whether a value is known before the graph runs.
Values that Nx cannot hold live in the graph environment as host terms: a string tensor is {:strings, shape, values}, a sequence is {:sequence, elements}, an empty optional is :none and a zero-size tensor is {:empty, shape}.
The shape and the type of the tensor that tensor/2 gives for t, from its dims and data_type, without its data. A string or an empty tensor gives nil.
Calls fun with values, where f16 and bf16 tensors become f32, and casts the f32 tensors that fun returns back to that type. fun takes and returns a list.
Sums, means and softmax in f16 or bf16 lose too much precision, so the Nx operations of such ops run in f32.