# `EMLX.Quantization`
[🔗](https://github.com/elixir-nx/emlx/blob/v0.4.2/emlx/lib/emlx/quantization.ex#L1)

Affine group-wise int2/int4/int8 quantization for Apple Silicon inference,
plus microscaled floating-point modes (`"mxfp4"`/`"mxfp8"`/`"nvfp4"`, via
`:mode` on `quantize/2` — see `EMLX.quantize/2`).

Quantized weights are represented as annotated `Nx.Tensor` values — the
tensor carries the original logical shape and type (e.g. `{:s, 4}` for
4-bit), while the `EMLX.Backend` struct stores the packed uint32 data and
a `EMLX.Quantization.Config` with scales, biases, group_size, bits, and
mode. Microscaled modes have no biases (`Config.biases` is `nil` — MLX's
`fp_quantize` returns only `(wq, scales)` for them).

`Nx.dot` automatically dispatches to `mx::quantized_matmul` when it detects
a quantized operand on `EMLX.Backend` — no explicit call site changes needed.

## Basic usage

    # Quantize a dense weight
    weight = Nx.iota({512, 4096}, type: :f32)
    weight = Nx.backend_transfer(weight, {EMLX.Backend, device: :gpu})
    qw = EMLX.Quantization.from_dense(weight)

    # Standard Nx.dot dispatches to mx::quantized_matmul automatically
    input = Nx.iota({1, 8, 4096}, type: :f32)
    result = Nx.dot(input, [2], qw, [1])

## Inside defn

`quantize/2`, `dequantize/1`, and `quantized_matmul/2` are all
`deftransform` functions backed by `Nx.runtime_call`, so they are safe
to call inside `Nx.Defn.jit`-traced forward passes:

    defn my_layer(x, qw) do
      dense = EMLX.Quantization.dequantize(qw)
      Nx.dot(x, [2], dense, [1])
    end

`Nx.dot` with a quantized tensor also dispatches transparently via
`EMLX.Backend.dot/7` at execution time, so explicit `dequantize` is
only needed when you want the dense weight for a non-dot operation.

## Limitation: freshly-quantized values can't feed a traced `Nx.dot`

`Nx.dot`'s `mx::quantized_matmul` dispatch (see
`EMLX.Native.Expr.quantizable_param_positions/1`) only recognizes a
quantized tensor that started life as a bound top-level parameter of the
`Nx.Defn.jit`-compiled function — it's resolved via a call-time
`quant_signature` keyed by parameter position, baked into a specialized
compiled program. A value produced by `quantize/2` *inside* the traced
program (e.g. `dequantize/1` → some op → `quantize/2` again) has no such
position and isn't recognized as a quantized `:dot` operand; it may only
be used as a direct output of the traced function (see `quantize/2`'s doc).

## See also

* `EMLX.Quantization.Config` — internal metadata struct

# `dequantize`

Dequantize a quantized `Nx.Tensor` via `Nx.runtime_call`.

At execution time the callback receives the real tensor, unpacks
`quantization_config`, and calls `mx::dequantize`. Safe to call inside
`Nx.Defn.jit`-traced forward passes.

The output has the same shape as the input (the quantized tensor's logical
shape equals the dense shape). Supports every mode `EMLX.quantize/2`
accepts.

# `gather_quantized_matmul`

Run a gathered quantized matmul via `Nx.runtime_call`.

The mixture-of-experts counterpart of `quantized_matmul/2`: `qw` stacks every
expert on its leading axis and `rhs_indices` names the expert for each row,
so the gather happens inside `mx::gather_qmm` and no per-token copy of the
expert weights is ever materialised. Safe inside `Nx.Defn.jit`-traced
forward passes.

Output shape is `{batch_dims..., rows, out_features}`, where the batch dims
come from broadcasting the activation's leading axes against `rhs_indices`
and `out_features` is the second dimension of `qw`. Output type follows the
same rule as `quantized_matmul/2`.

## Options

  * `:lhs_indices` - reorders the activation rows; `nil` takes them in order
  * `:sorted_indices` - set when `rhs_indices` is already sorted, so the
    kernel can skip its own sort

# `quantize`

Quantize a dense tensor of rank 2 or higher via `Nx.runtime_call`.

Quantization runs along the last axis; leading axes are batch or stacking
dims.

At execution time the callback receives the real tensor and runs
`mx::quantize`. Returns an annotated quantized `Nx.Tensor` with the same
logical shape as the input and type `{:s, N}`. Safe to call inside
`Nx.Defn.jit`-traced forward passes — including on a value that is itself
freshly computed in-graph (e.g. `dequantize/1` followed by some op followed
by `quantize/2` again), not just a bare top-level parameter.

Note: the *result* of a traced `quantize/2` call may only be used as a
direct output of the `Nx.Defn.jit`-compiled function (returned as-is, or
passed through unchanged). Feeding it into another op in the same traced
program (e.g. `Nx.dot`) is not yet supported and raises during lowering —
only a quantized tensor that started life as a bound top-level parameter
dispatches to `mx::quantized_matmul` today.

## Options

* `:type` — Nx storage type: `{:s, 2}`, `{:s, 4}` (default), or `{:s, 8}`.
* `:group_size` — 32, 64, or 128 (default 64). Must evenly divide the last
  dimension of `tensor`.
* `:mode` — `"affine"` (default), or a microscaled mode (`"mxfp4"`,
  `"mxfp8"`, `"nvfp4"`) — see `EMLX.quantize/2`.

# `quantized?`

Returns `true` if the tensor has quantization metadata on its backend.

# `quantized_matmul`

Run a quantized matmul via `Nx.runtime_call`.

At execution time the callback unpacks `quantization_config` from `qw` and
calls `mx::quantized_matmul`. Safe to call inside `Nx.Defn.jit`-traced
forward passes.

Output shape is `{batch_dims..., out_features}` where `out_features` is the
first dimension of `qw`. Output type matches the scales dtype for
`"affine"` (historically the activation's own float type); for microscaled
modes scales are `:u8` (a packed exponent byte), so the output type
follows the activation's dtype instead.

# `quantized_tensor`

Construct an annotated quantized `Nx.Tensor` from pre-computed device refs.

Use this when you already have packed weights from a checkpoint. For
quantizing a dense tensor from scratch, prefer `from_dense/2`.

## Options

* `:type` — Nx storage type: `{:s, 2}`, `{:s, 4}` (default), or `{:s, 8}`.
* `:group_size` — quantization group size (default 64).

---

*Consult [api-reference.md](api-reference.md) for complete listing*