Feature operations¶
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.
Related pages¶
Sparse tensor identity: Sparse tensor model
Feature module wrappers: Feature modules
Quantized weight storage: Quantized weights
Quantized route details: Quantization routes
- 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
meanorvaris omitted, statistics are computed from active feature rows. Optional affineweightandbiashave shape(C,).- Return type:
- Parameters:
x (SparseTensor)
weight (array | None)
bias (array | None)
mean (array | None)
var (array | None)
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
trainingis false orpis zero, the feature matrix is returned unchanged inside a new sparse wrapper.- Return type:
- Parameters:
x (SparseTensor)
p (float)
training (bool)
- mlx_lattice.ops.feature.gelu(x, *, approximate='none')[source]¶
Apply GELU to sparse features while preserving coordinates.
approximateaccepts'none'/'precise'for the erf formula,'tanh'for the tanh approximation, or'fast'for the sigmoid-based approximation.- Return type:
- Parameters:
x (SparseTensor)
approximate (Literal['none', 'precise', 'tanh', 'fast'])
- 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:
- Parameters:
x (SparseTensor)
weight (array | None)
bias (array | None)
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:
- Parameters:
x (SparseTensor)
negative_slope (float)
- 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 useQuantizedWeightwithlinearlayout. Optional bias has shape(C_out,).- Return type:
- Parameters:
x (SparseTensor)
weight (array | QuantizedWeight)
bias (array | None)
- mlx_lattice.ops.feature.relu(x)[source]¶
Apply ReLU to sparse features while preserving coordinates.
- Return type:
- 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:
- Parameters:
x (SparseTensor)
weight (array | None)
eps (float)
- mlx_lattice.ops.feature.sigmoid(x)[source]¶
Apply sigmoid to sparse features while preserving coordinates.
- Return type:
- Parameters:
x (SparseTensor)
- mlx_lattice.ops.feature.silu(x)[source]¶
Apply SiLU/Swish to sparse features while preserving coordinates.
- Return type:
- 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
thresholdin the scaled domain return the input directly to avoid unnecessary exponential work.- Return type:
- Parameters:
x (SparseTensor)
beta (float)
threshold (float)
- mlx_lattice.ops.feature.tanh(x)[source]¶
Apply hyperbolic tangent to sparse features while preserving coordinates.
- Return type:
- Parameters:
x (SparseTensor)