Quantization

Post-training weight quantization utilities and Qwix integration.

taktiny.utils.quantization.quantize_linear_weight(array, rule, input_axis_count, batch_axis_count=0, scale_dtype=None)

Return type:

Any

Parameters:

• array (Any)

• rule (Any)

• input_axis_count (int)

• batch_axis_count (int)

• scale_dtype (Any)

taktiny.utils.quantization.quantize_conv_weight(array, rule, output_axis_count=1, scale_dtype=None)

Quantizes a convolution kernel per structured output channel.

Return type:

Any

Parameters:

• array (Any)

• rule (Any)

• output_axis_count (int)

• scale_dtype (Any)

taktiny.utils.quantization.quantize_embedding_weight(array, rule, scale_dtype=None, vocabulary_axis_count=1)

Return type:

Any

Parameters:

• array (Any)

• rule (Any)

• scale_dtype (Any)

• vocabulary_axis_count (int)

taktiny.utils.quantization.quantization_rules(quantization)

Return type:

tuple[QuantizationRule, ...]

Parameters:

quantization (Any)

taktiny.utils.quantization.quantize_model(model, provider, *model_inputs, methods=('__call__',), **model_inputs_kwargs)

Return an independent model using Qwix quantized training operations.

Supports the standard qwix.QtProvider with dynamic activation scales for lax.dot_general, two-operand jnp.einsum, and lax.conv_general_dilated. Parameters remain floating point; Qwix quantizes operation operands and, when bwd_qtype is set, backward matrix products. This is not conversion to permanently integer weights.

Rules use first-match precedence and dotted or slash-separated module paths (the root path is an empty string). Selected entry methods are intercepted, including operations in their child modules. Calls through other root methods are unchanged. Existing custom operation hooks must ultimately call one of the supported JAX operations to be intercepted. Shared child modules use their first traversal path for rule matching. Weight detection for dot products uses parameter-array identity; if a custom module transforms a weight before the dot, it is treated as an activation. Convolutions always treat the right operand as a weight.

Each selected method is run once on a disposable copy with the supplied inputs to validate execution without changing the returned model’s state. The same inputs must be valid for every selected method. Python branches not exercised by these inputs are only checked when subsequently called.

Parameters:

• model (Module) – A Taktiny module with floating-point parameters.

• provider (QuantizationProvider) – A standard Qwix QtProvider. Other providers, static scales, and stochastic rounding are not supported yet and raise errors.

• *model_inputs (Any) – Positional example inputs for the selected methods.

• methods (Collection[str]) – Root entry methods to intercept. Defaults to ('__call__',).

• **model_inputs_kwargs (Any) – Keyword example inputs, forwarded unchanged.

Return type:

Module

Returns:

A quantized model instance.

Example

>>> import jax.numpy as jnp
>>> import qwix
>>> from taktiny import nn
>>> model = nn.Linear(8, 4, rngs=nn.Rngs(0))
>>> provider = qwix.QtProvider([qwix.QtRule(
...     weight_qtype='int8', act_qtype='int8', bwd_qtype='int8')])
>>> model = quantize_model(model, provider, jnp.ones((2, 8)))
>>> model(jnp.ones((2, 8))).shape
(2, 4)

Note

Uses private Qwix kernel/interception APIs; Qwix upgrades should be checked against the integration tests. Precision hints are handled by the QT kernels, not by ordinary floating-point dot precision settings.