Feature operations
Sparse feature operations preserve sparse coordinate identity. They replace
only the feats matrix and keep the coordinate manager, coordinate key,
stride, active rows, and batch metadata unchanged.
Use these functions when the operation is row-local or channel-local:
activations, normalization, dropout, and linear projections. A quantized linear
projection is selected by passing mlx_lattice.core.QuantizedWeight.
The MLIR artifact ABI separates dense feature math from sparse identity.
linear_features, activation, batch_norm_features,
layer_norm_features, and rms_norm_features are dense rank-2 feature
tensor operations. Sparse wrappers such as relu, linear,
batch_norm, layer_norm, and rms_norm apply those dense ops to
x.feats and return x.replace(feats=...). Serialized MLIR graphs must
represent the same pattern with lattice.sparse.decompose and
lattice.sparse.with_features.
Related pages
-
mlx_lattice.ops.feature.activation(x, *, kind, approximate='none', alpha=0.01, beta=1.0, threshold=20.0)[source]
Apply a dense feature activation to a rank-2 feature tensor.
- Return type:
array
- Parameters:
x (array)
kind (Literal['relu', 'sigmoid', 'gelu', 'silu', 'leaky_relu', 'tanh', 'softplus'])
approximate (Literal['none', 'precise', 'tanh', 'fast'])
alpha (float)
beta (float)
threshold (float)
-
mlx_lattice.ops.feature.batch_norm(x, *, weight=None, bias=None, mean=None, var=None, eps=1e-05)[source]
Apply per-channel batch normalization to sparse features.
If mean or var is omitted, statistics are computed from active
feature rows. Optional affine weight and bias have shape (C,).
- Return type:
SparseTensor
- Parameters:
x (SparseTensor)
weight (array | None)
bias (array | None)
mean (array | None)
var (array | None)
eps (float)
-
mlx_lattice.ops.feature.batch_norm_features(x, scale, bias, mean, var, *, eps=1e-05)[source]
Apply frozen-stat batch normalization to rank-2 feature rows.
- Return type:
array
- Parameters:
x (array)
scale (array)
bias (array)
mean (array)
var (array)
eps (float)
-
mlx_lattice.ops.feature.dropout(x, *, p=0.5, training=True)[source]
Apply inverted dropout to sparse features during training.
Coordinates are preserved. When training is false or p is zero, the
feature matrix is returned unchanged inside a new sparse wrapper.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.feature.gelu(x, *, approximate='none')[source]
Apply GELU to sparse features while preserving coordinates.
approximate accepts 'none'/'precise' for the erf formula,
'tanh' for the tanh approximation, or 'fast' for the sigmoid-based
approximation.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.feature.layer_norm(x, *, weight=None, bias=None, eps=1e-05)[source]
Apply layer normalization independently to each sparse row.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.feature.layer_norm_features(x, scale, bias, *, eps=1e-05)[source]
Apply layer normalization to rank-2 feature rows.
- Return type:
array
- Parameters:
x (array)
scale (array)
bias (array)
eps (float)
-
mlx_lattice.ops.feature.leaky_relu(x, *, negative_slope=0.01)[source]
Apply leaky ReLU to sparse features while preserving coordinates.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.feature.linear(x, weight, bias=None)[source]
Apply a dense or quantized linear projection to sparse features.
Coordinates are preserved. Dense weights use shape (C_out, C_in) and
packed weights use QuantizedWeight with linear layout. Optional
bias has shape (C_out,).
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.feature.linear_features(x, weight, bias=None)[source]
Apply a dense or quantized linear projection to a feature matrix.
x must have shape (N, C_in). Dense weights use shape
(C_out, C_in) and packed weights use QuantizedWeight with
linear layout. Optional bias has shape (C_out,).
- Return type:
array
- Parameters:
-
-
mlx_lattice.ops.feature.relu(x)[source]
Apply ReLU to sparse features while preserving coordinates.
- Return type:
SparseTensor
- Parameters:
x (SparseTensor)
-
mlx_lattice.ops.feature.rms_norm(x, *, weight=None, eps=1e-05)[source]
Apply RMS normalization independently to each sparse row.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.feature.rms_norm_features(x, scale, *, eps=1e-05)[source]
Apply RMS normalization to rank-2 feature rows.
- Return type:
array
- Parameters:
x (array)
scale (array)
eps (float)
-
mlx_lattice.ops.feature.sigmoid(x)[source]
Apply sigmoid to sparse features while preserving coordinates.
- Return type:
SparseTensor
- Parameters:
x (SparseTensor)
-
mlx_lattice.ops.feature.silu(x)[source]
Apply SiLU/Swish to sparse features while preserving coordinates.
- Return type:
SparseTensor
- Parameters:
x (SparseTensor)
-
mlx_lattice.ops.feature.softplus(x, *, beta=1.0, threshold=20.0)[source]
Apply numerically thresholded softplus to sparse features.
Values above threshold in the scaled domain return the input directly
to avoid unnecessary exponential work.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.feature.tanh(x)[source]
Apply hyperbolic tangent to sparse features while preserving coordinates.
- Return type:
SparseTensor
- Parameters:
x (SparseTensor)