Module ops 

Primitive tensor operations: f32 only, no allocations.

Functions

add_bias

Add bias in-place: x[b, j] += bias[j]

dot

Dot product of two equal-length slices.

elementwise_mul

Element-wise multiply: y[i] *= x[i]

gelu

GELU activation (approximate, matches JAX nn.gelu default). Formula: 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))

matmul

Dense matmul: y[b,j] = Σ_i x[b,i] * w[i,j] w shape [in_dim, out_dim], x shape [batch, in_dim], y shape [batch, out_dim].

saxpy

y += alpha * x (SAXPY)

sigmoid

Element-wise sigmoid: σ(x) = 1 / (1 + e^-x).

silu

SiLU / swish: x * sigmoid(x).

softmax_inplace

In-place softmax over a slice (stable via max subtraction).