Linear Layers¶

Standard and parameter-efficient linear layers.

class taktiny.nn.Linear(in_features, out_features, *, bias=True, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: Module

Applies a linear transformation to the input.

\[Y = XW + b\]

Here, W is the learnable kernel and b is the optional learnable bias. For one-dimensional features, an input of shape (..., in_features) produces an output of shape (..., out_features).

When feature shapes are tuples, the same operation contracts all trailing input-feature axes: a kernel of shape (*in_features, *out_features) maps (..., *in_features) to (..., *out_features). The bias term is omitted when bias=False.

Parameters:

in_features (int | Sequence[int]) – Size of the trailing input feature axes. An integer is treated as a one-dimensional feature shape.
out_features (int | Sequence[int]) – Size of the output feature axes. An integer is treated as a one-dimensional feature shape.
bias (bool) – Whether to create and add a learnable bias. Defaults to True.
dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to the parameter initializers. Defaults to the initializer’s default data type.
rngs (Rngs) – Random number generator used to initialize the kernel and bias.
kernel_initializer (Initializer) – Function used to initialize the kernel. Defaults to LeCun uniform initialization.
bias_initializer (Initializer) – Function used to initialize the bias. Defaults to zeros.
quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for the kernel.
dot_general (DotGeneral | None) – Optional implementation of dot_general. It is used for non-quantized kernels instead of jax.lax.dot_general.
axis_names (tuple[str | None, ...] | None) – Optional logical axis names for the kernel. The names of the output axes are also assigned to the bias.
partition_spec (P | None) – Optional partition specification for the kernel. The specification for the output axes is also assigned to the bias.
kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to the kernel parameter.
bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to the bias parameter.
precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Dot-product precision forwarded to dot_general.
preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred result and accumulation data type forwarded to dot_general.

Variables:

kernel – The learnable kernel parameter.
bias – The learnable bias parameter, or None when bias is disabled.

Examples

Apply a conventional projection to the last axis:

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> linear = nn.Linear(3, 2, rngs=nn.Rngs(0))
>>> x = jnp.ones((4, 3))
>>> linear(x).shape
(4, 2)

Contract multiple input axes and produce multiple output axes:

>>> linear = nn.Linear((2, 3), (4, 5), rngs=nn.Rngs(1))
>>> x = jnp.ones((8, 2, 3))
>>> linear(x).shape
(8, 4, 5)

class taktiny.nn.Bilinear(in1_features, in2_features, out_features, *, bias=True, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: Module

Applies a bilinear transformation to two inputs.

\[y_k = x_1^{\mathsf{T}} W_k x_2 + b_k\]

Here, W_k is the learnable kernel for output feature k and b_k is the optional bias. For one-dimensional features, inputs of shape (..., in1_features) and (..., in2_features) produce an output of shape (..., out_features).

When feature shapes are tuples, the same operation contracts all trailing feature axes. The kernel has shape (*in1_features, *in2_features, *out_features).

Parameters:

in1_features (int | Sequence[int]) – Size of the trailing feature axes of the first input. An integer is treated as a one-dimensional feature shape.
in2_features (int | Sequence[int]) – Size of the trailing feature axes of the second input. An integer is treated as a one-dimensional feature shape.
out_features (int | Sequence[int]) – Size of the output feature axes. An integer is treated as a one-dimensional feature shape.
bias (bool) – Whether to create and add a learnable bias. Defaults to True.
dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to the parameter initializers. Defaults to the initializer’s default data type.
rngs (Rngs) – Random number generator used to initialize the kernel and bias.
kernel_initializer (Initializer) – Function used to initialize the kernel. Defaults to LeCun uniform initialization.
bias_initializer (Initializer) – Function used to initialize the bias. Defaults to zeros.
quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for the kernel.
dot_general (DotGeneral | None) – Optional implementation of dot_general. It is used for non-quantized kernels instead of jax.lax.dot_general.
axis_names (tuple[str | None, ...] | None) – Optional logical axis names for the kernel. The names of the output axes are also assigned to the bias.
partition_spec (P | None) – Optional partition specification for the kernel. The specification for the output axes is also assigned to the bias.
kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to the kernel parameter.
bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to the bias parameter.
precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Dot-product precision forwarded to dot_general.
preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred result and accumulation data type forwarded to dot_general.

Variables:

kernel – The learnable bilinear kernel parameter.
bias – The learnable bias parameter, or None when bias is disabled.

Examples

Apply a bilinear transformation to two batches of vectors:

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> bilinear = nn.Bilinear(3, 4, 2, rngs=nn.Rngs(0))
>>> x1 = jnp.ones((5, 3))
>>> x2 = jnp.ones((5, 4))
>>> bilinear(x1, x2).shape
(5, 2)

Contract multi-axis feature shapes:

>>> bilinear = nn.Bilinear(
...     (2, 3), (4, 5), (6, 7), rngs=nn.Rngs(1)
... )
>>> x1 = jnp.ones((8, 2, 3))
>>> x2 = jnp.ones((8, 4, 5))
>>> bilinear(x1, x2).shape
(8, 6, 7)

class taktiny.nn.LoRALinear(base, rank, alpha, *, bias=None, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: Module

Add a low-rank update to an existing Linear module.

Computes base(x) + (alpha / rank) * lora_B(lora_A(x)). A contracts all input-feature axes into the rank dimension; B projects that dimension into all output-feature axes. B’s kernel starts at zero, so the default zero bias makes the initial result equal to base(x). The base is put in evaluation mode; this does not freeze its gradients. Select adapter parameters in the optimizer when training only LoRA.

Omitted (or None) dtype, bias, metadata, and dot settings inherit from the base. When both sharding arguments are None, the base’s logical labels and partition specification are inherited. Linear resolves those labels using the current mapping rules, overriding the inherited spec. Outside the mapping context, unmapped labels therefore become replicated. A base without logical names supplies its physical specification instead. Providing either sharding argument selects an adapter-specific layout. Use partition_spec=PartitionSpec() for replicated adapters. Initializers and quantization keep their adapter defaults because Linear does not retain their original constructor configurations.

Parameters:

base (Linear) – Existing Linear, including one with N-D feature shapes.
rank (int) – Positive adapter rank.
alpha (float) – Finite scaling numerator; the update is scaled by alpha/rank.
rngs (Rngs) – Random stream used to initialize adapter parameters.
bias (bool | None) – Create an output bias in B. A is always bias-free. The base bias is unaffected. None inherits whether the base has a bias. Set False for the usual bias-free LoRA update.
dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Adapter initializer dtype. None inherits the base kernel dtype (the scale dtype for a quantized base).
kernel_initializer (Initializer) – Initializer for A; B always uses zeros.
bias_initializer (Initializer) – Initializer for B’s optional bias. A nonzero bias makes the initial adapter update nonzero when alpha is nonzero.
quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix configuration for adapter kernels only. Defaults to dense adapters even when the base is quantized.
dot_general (DotGeneral | None) – Custom dot operation for non-quantized adapter kernels; None inherits the base implementation.
axis_names (tuple[str | None, ...] | None) – Logical names in base-kernel order: input axes followed by output axes. A and B receive their respective feature names and an unnamed rank axis.
partition_spec (P | None) – Explicit base-kernel specification, split between A and B with a replicated rank axis. Omitted trailing entries are replicated. Logical mappings override explicit specs.
kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to both adapter kernels. None copies the base kernel metadata; an empty dict clears it.
bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to B’s bias. None copies the base bias metadata, when present; an empty dict clears it.
precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Dot precision for both adapter projections; None inherits the base precision.
preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred adapter accumulation/result dtype; None inherits the base setting.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> base = nn.Linear((2, 3), (4, 5), rngs=nn.Rngs(0))
>>> layer = nn.LoRALinear(base, 2, 4.0, rngs=nn.Rngs(1))
>>> layer(jnp.ones((7, 2, 3))).shape
(7, 4, 5)

class taktiny.nn.DoRALinear(base, rank, alpha, *, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: _WeightAdapter

Adapt weight direction with LoRA and learn a separate output magnitude.

Computes a column-normalized update to the base kernel. The base bias is added after magnitude scaling. A uses kernel_initializer and B starts at zero. Magnitude starts at the base column norm, so initialization preserves the base output. Zero columns are handled with a finite norm floor.

Dtype, metadata, dot settings and sharding inherit as in LoRALinear. axis_names/partition_spec describe the base kernel’s input then output axes; the adapter rank is replicated and magnitude uses output axes. Adapter factors are bias-free and dense, including for a quantized base. Base gradients are not frozen automatically; select adapter parameters in the optimizer. The direction norm is detached during differentiation.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> base = nn.Linear((2, 3), 4, rngs=nn.Rngs(0))
>>> nn.DoRALinear(base, 2, 4.0, rngs=nn.Rngs(1))(jnp.ones((5, 2, 3))).shape
(5, 4)

Reference:: Liu et al., “DoRA: Weight-Decomposed Low-Rank Adaptation” (2024). https://arxiv.org/abs/2402.09353

Parameters:

base (Linear)
rank (int)
alpha (float)
dtype (DType | None)
rngs (Rngs)
kernel_initializer (Initializer)
dot_general (DotGeneral | None)
axis_names (AxisNames | None)
partition_spec (PartitionSpec | None)
kernel_metadata (MetaData | None)
precision (PrecisionLike)
preferred_element_type (DTypeLike | None)

class taktiny.nn.AdaLoRALinear(base, rank, alpha, *, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: _WeightAdapter

An SVD-style adapter: base(x) + alpha/rank * B(E * A(x)).

A and B use kernel_initializer; the replicated rank vector E starts at zero. mask_rank performs one pruning update. Budget scheduling, importance estimation and adding orthogonal_loss to a training objective are the caller’s responsibility.

Dtype, kernel metadata, dot settings and sharding inherit as in LoRALinear. Logical axes are resolved in the current context; the rank stays replicated. No adapter biases are created. A quantized base may be wrapped, while the new factors remain floating point.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> layer = nn.AdaLoRALinear(nn.Linear(3, 4, rngs=nn.Rngs(0)), 2, 4.0, rngs=nn.Rngs(1))
>>> layer(jnp.ones((5, 3))).shape
(5, 4)

Reference:: Zhang et al., “AdaLoRA: Adaptive Budget Allocation for Parameter-Efficient Fine-Tuning” (2023). https://arxiv.org/abs/2303.10512

Parameters:

base (Linear)
rank (int)
alpha (float)
dtype (DType | None)
rngs (Rngs)
kernel_initializer (Initializer)
dot_general (DotGeneral | None)
axis_names (AxisNames | None)
partition_spec (PartitionSpec | None)
kernel_metadata (MetaData | None)
precision (PrecisionLike)
preferred_element_type (DTypeLike | None)

class taktiny.nn.LoHaLinear(base, rank, alpha, *, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: _WeightAdapter

Hadamard-product adapter with delta W = (A1 B1) * (A2 B2).

The forward pass uses an effective rank of rank squared, avoiding a full input-by-output update kernel. B2 starts at zero; the other factors use kernel_initializer. The base bias is unchanged.

All factors support N-D feature shapes, inherited dtype and metadata, and base-kernel logical/physical specifications with a replicated rank axis. Current logical rules take precedence, as in Linear. Quantized bases are supported with dense adapter factors.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> layer = nn.LoHaLinear(nn.Linear(3, 4, rngs=nn.Rngs(0)), 2, 4.0, rngs=nn.Rngs(1))
>>> layer(jnp.ones((5, 3))).shape
(5, 4)

Reference:: Yeh et al., “Navigating Text-To-Image Customization: From LyCORIS Fine-Tuning to Model Evaluation” (2024). https://arxiv.org/abs/2309.14859

Parameters:

base (Linear)
rank (int)
alpha (float)
dtype (DType | None)
rngs (Rngs)
kernel_initializer (Initializer)
dot_general (DotGeneral | None)
axis_names (AxisNames | None)
partition_spec (PartitionSpec | None)
kernel_metadata (MetaData | None)
precision (PrecisionLike)
preferred_element_type (DTypeLike | None)

class taktiny.nn.LoKrLinear(base, rank, alpha, *, dtype=None, rngs, decompose_both=False, decompose_factor=-1, kernel_initializer=<function variance_scaling.<locals>.init>, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: _WeightAdapter

Kronecker-product weight adapter, delta W = kron(W1, W2).

decompose_factor controls the split of flattened input/output widths; -1 selects balanced factors. Small ranks factorize W2, and decompose_both=True also permits factorizing W1. W1 (or its A factor) starts at zero; other factors use kernel_initializer.

Dtype, metadata, dot settings, and the effective kernel’s sharding inherit from the base. Kronecker factors are replicated because flattened factor dimensions do not correspond to the original N-D feature axes. A sharded effective kernel is materialized when a non-replicated layout is requested. Otherwise the forward pass contracts the two small factors directly. Base biases are unchanged; adapter factors stay dense for quantized bases.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> layer = nn.LoKrLinear(nn.Linear((2, 3), 4, rngs=nn.Rngs(0)), 2, 4.0, rngs=nn.Rngs(1))
>>> layer(jnp.ones((5, 2, 3))).shape
(5, 4)

Reference:: Yeh et al., “Navigating Text-To-Image Customization: From LyCORIS Fine-Tuning to Model Evaluation” (2024). https://arxiv.org/abs/2309.14859

Parameters:

base (Linear)
rank (int)
alpha (float)
dtype (DType | None)
rngs (Rngs)
decompose_both (bool)
decompose_factor (int)
kernel_initializer (Initializer)
dot_general (DotGeneral | None)
axis_names (AxisNames | None)
partition_spec (PartitionSpec | None)
kernel_metadata (MetaData | None)
precision (PrecisionLike)
preferred_element_type (DTypeLike | None)

class taktiny.nn.VeRALinear(base, rank, *, vera_A, vera_B, dtype=None, d_initial=0.1, dot_general=None, axis_names=None, partition_spec=None, metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: _WeightAdapter

Learn rank/output scales around shared, fixed random projections.

Computes base(x) + ((x A) * lambda_d) B * lambda_b. Supplied A and B are two-dimensional Parameters and may be larger than this layer: prefixes matching flattened feature widths and rank are selected. Their values, metadata, and sharding are not modified. stop_gradient makes them fixed in differentiation; exclude them from optimizer weight decay as well.

lambda_d is replicated and initialized to d_initial (default 0.1). lambda_b starts at zero and uses the base output feature shape and axes. Dtype, metadata, dot settings and logical/physical axes inherit from base. Logical names are resolved in the current context. No RNG is needed: randomness is entirely in the supplied projections. Quantized bases and supplied Qwix projections are read in floating point for the adapter.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> base = nn.Linear(3, 4, rngs=nn.Rngs(0))
>>> a = nn.Parameter(jnp.ones((3, 2)), trainable=False)
>>> b = nn.Parameter(jnp.ones((2, 4)), trainable=False)
>>> layer = nn.VeRALinear(base, 2, vera_A=a, vera_B=b)
>>> layer(jnp.ones((5, 3))).shape
(5, 4)

Reference:: Kopiczko et al., “VeRA: Vector-based Random Matrix Adaptation” (2024). https://arxiv.org/abs/2310.11454

Parameters:

base (Linear)
rank (int)
vera_A (Parameter)
vera_B (Parameter)
dtype (DType | None)
d_initial (float)
dot_general (DotGeneral | None)
axis_names (AxisNames | None)
partition_spec (PartitionSpec | None)
metadata (MetaData | None)
precision (PrecisionLike)
preferred_element_type (DTypeLike | None)