Operations

Layout and quantization aware high-level math operations.

taktiny.utils.ops.linear(x, w, b=None, *, quant=None, module_path='', precision=None, preferred_element_type=None, out_sharding=None)[source]

Compute x @ w + b with optional quantized multiplication.

x has shape (..., input_features) and w has shape (input_features, output_features), matching Taktiny Linear’s kernel layout. Leading input dimensions are preserved. Bias must broadcast to the result; a quantized bias is dequantized before addition.

quant=None uses JAX for dense inputs and Qwix for QArray inputs. quant accepts a dtype string, a Qwix rule, a sequence of rules, or a PtqProvider/QtProvider. Strings use Taktiny’s weight-only shorthand. Rules match module_path (empty by default) and dot_general in first-match order. weight_qtype applies to w and act_qtype to x. Existing QArrays are dequantized before applying a selected rule. A rule with unset qtypes disables quantization for its matched operation.

QtRule uses Qwix’s custom-gradient training kernels; bwd_qtype controls backward quantization. Ordinary rules and string shortcuts use PTQ, not trainable weight quantization. For training, pass floating-point parameters and a QtRule. Static activation scales and stochastic rounding are not supported by these stateless functions. QT currently does not support an explicit out_sharding argument. Training uses Qwix QT kernels; supported formats and backward operations depend on Qwix and the execution backend.

Returns a dense JAX array. preferred_element_type selects its dtype, including after bias addition; otherwise the dtype is promoted from the inputs and bias, using QArray scale dtypes. Raw FP8 inputs promote to BF16; integer-only inputs with quant set promote to FP32. Accumulation may use a higher precision than the returned dtype. out_sharding is passed to the dot operation where supported.

Return type:

Array

Parameters:

Example

>>> linear(jnp.ones((2, 3)), jnp.ones((3, 4))).shape
(2, 4)
>>> linear(jnp.ones((2, 3)), jnp.ones((3, 4)), quant='int8',
...        preferred_element_type=jnp.bfloat16).dtype
dtype(bfloat16)
taktiny.utils.ops.einsum(subscripts, *operands, quant=None, module_path='', optimize='auto', precision=None, preferred_element_type=None, out_sharding=None)[source]

Evaluate an Einstein summation with optional quantized contractions.

Uses the string-equation form of jax.numpy.einsum. Dense array-like inputs are converted to JAX arrays, while QArrays retain their quantized representation when quant and out_sharding are unset. In that case, QArray inputs select qwix.einsum and its dequantization fallbacks. quant accepts the same rule configurations as linear(). Strings are weight-only shortcuts. In two-operand expressions the first operand is the activation and the second is the weight, independent of contraction ordering. Rules match module_path and either einsum or dot_general. QtRule enables training and its bwd_qtype controls backward quantization. Rule-based expressions support at most two operands; transformations that obscure operand roles raise an error. Unary operations remain ordinary JAX operations. Returns a dense JAX array.

Static scales, stochastic rounding, and explicit output sharding for QT have the same limitations as linear().

precision is forwarded to the selected backend. optimize is supported for dense inputs and when quant is set. With quant=None, unsharded QArray inputs require 'auto' because Qwix manages its own contraction path. preferred_element_type selects the returned dtype; otherwise input dtypes (QArray scale dtypes) are promoted. Raw FP8 promotes to BF16, and integer-only inputs promote to FP32 when quant is set. Accumulation may use higher precision. out_sharding is passed to JAX’s einsum. If quant is unset, this sharded path dequantizes QArrays and uses dense computation.

Return type:

Array

Parameters:

Example

>>> einsum('...i,io->...o', jnp.ones((2, 3)), jnp.ones((3, 4))).shape
(2, 4)
taktiny.utils.ops.conv_general_dilated(lhs, rhs, window_strides, padding, lhs_dilation=None, rhs_dilation=None, dimension_numbers=None, feature_group_count=1, batch_group_count=1, precision=None, preferred_element_type=None, *, out_sharding=None, quant=None, module_path='')[source]

Convolve activations lhs with kernel rhs using JAX or Qwix.

Layout, strides, dilation, and grouping follow jax.lax.conv_general_dilated. dimension_numbers describes input, kernel, and output layouts, e.g. ('NWC', 'WIO', 'NWC') for channels-last 1-D convolution. Padding is 'VALID', 'SAME', 'SAME_LOWER', or one (low, high) pair per spatial dimension. Dilation inserts gaps in the input or kernel; strides control output sampling.

Dense array-like inputs use JAX; QArrays use Qwix. quant accepts a weight-only dtype string, a Qwix rule, a rule sequence, or PtqProvider / QtProvider. Rules match module_path and conv_general_dilated in first-match order. A selected rule dequantizes existing QArrays before applying its act_qtype to lhs and weight_qtype to rhs. Unmatched or disabled rules use dense computation.

Ordinary rules quantize operands at each call. For training, use floating- point parameters and QtRule; bwd_qtype controls backward quantization through Qwix QT. Supported formats and gradients depend on Qwix and the backend. Tiled convolution quantization, static activation calibration, QT stochastic rounding, and convolution additional_qt_config are unsupported.

Returns a dense array. preferred_element_type selects the result dtype; otherwise operand dtypes (scale dtypes for QArrays) are promoted. FP8 inputs promote to BF16 and quantized integer inputs to FP32. Precision is forwarded to JAX/PTQ; Qwix QT controls its own contraction precision. out_sharding is forwarded to the convolution operation.

Return type:

Array

Parameters:

Example

>>> conv_general_dilated(
...     jnp.ones((1, 8, 2)), jnp.ones((3, 2, 4)), (1,), 'SAME',
...     dimension_numbers=('NWC', 'WIO', 'NWC'),
... ).shape
(1, 8, 4)
taktiny.utils.ops.dot_general(lhs, rhs, dimension_numbers, precision=None, preferred_element_type=None, *, out_sharding=None, quant=None, module_path='', lhs_weight=False, rhs_weight=True)[source]

Contract two operands using JAX or Qwix dot_general.

dimension_numbers has the same contracting and batch-axis layout as jax.lax.dot_general(). Dense inputs use JAX; without an explicit rule, QArray inputs use Qwix. A selected quant rule dequantizes any existing QArrays before applying its configured operand quantization. Rules match module_path and dot_general in first-match order.

By default, lhs is an activation and rhs is a weight. Set lhs_weight or rhs_weight to choose which operand receives the rule’s weight_qtype; other operands receive act_qtype. A qwix.QtRule uses Qwix’s training kernel, including its optional backward quantization. Ordinary rules and strings use PTQ. QT does not currently support explicit out_sharding.

The result is a dense JAX array. Its implicit dtype follows the operand dtypes (QArray scale dtypes), with FP8 promoted to BF16 and quantized integer inputs to FP32. preferred_element_type overrides the result dtype. out_sharding is forwarded to the selected operation where supported.

Return type:

Array

Parameters:

Example

>>> dot_general(
...     jnp.ones((2, 3)), jnp.ones((3, 4)),
...     (((1,), (0,)), ((), ())),
... ).shape
(2, 4)