lora_matmul

Contents

lora_matmul#

class braintrace.lora_matmul(x, B, A, *, alpha=1.0, bias=None, b_fn=None, a_fn=None, bias_fn=None)[source]#

ETP-aware LoRA (Low-Rank Adaptation) matrix multiplication.

Computes \(y = \alpha \cdot x \mathbin{@} b\_fn(B) \mathbin{@} a\_fn(A) \; (+ bias\_fn(b))\), routing both low-rank factors (and the optional bias) through an ETP primitive so they participate in eligibility-trace computation. Auto-dispatches batched/unbatched based on x.ndim.

Parameters:

  • x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Input array, shape (batch, in_features) or (in_features,). Higher-rank x (x.ndim > 2) is rejected with a ValueError: every ETP trace rule assumes one of these two layouts.

  • B (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Low-rank matrix \(B\), shape (in_features, rank).

  • A (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Low-rank matrix \(A\), shape (rank, out_features).

  • alpha (float) – Scalar scaling factor \(\alpha\). Default 1.0.

  • bias (Array | ndarray | bool | number | bool | int | float | complex | Quantity | None) – Bias vector, shape (out_features,). Default None.

  • b_fn (Callable[[Array | ndarray | bool | number | bool | int | float | complex | Quantity], Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) – Elementwise transform applied to the B factor before the matrix multiplication. b_fn(B) must return an array of the same shape as B. None means identity (no transform). The VJP of b_fn is auto-composed inside xy_to_dw so that gradients w.r.t. the raw lora_b weights are correct. The transform operates on the unitless mantissa; physical units are split off before and recombined after. Pass a module-level function, not a fresh lambda, if this is called repeatedly: the hook is stored as a static eqn.params entry hashed by object identity, so two textually identical lambda objects are cache misses and silently retrace every call.

  • a_fn (Callable[[Array | ndarray | bool | number | bool | int | float | complex | Quantity], Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) – Elementwise transform applied to the A factor before the matrix multiplication. a_fn(A) must return an array of the same shape as A. None means identity. The transform operates on the unitless mantissa; physical units are split off before and recombined after. Same re-tracing caveat as b_fn: pass a module-level function rather than a fresh lambda.

  • bias_fn (Callable[[Array | ndarray | bool | number | bool | int | float | complex | Quantity], Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) – Elementwise transform applied to bias before adding. None means identity. The transform operates on the unitless mantissa; physical units are split off before and recombined after. Same re-tracing caveat as b_fn: pass a module-level function rather than a fresh lambda.

Returns:

Output array, shape (batch, out_features) or (out_features,).

Return type:

Array | ndarray | bool | number | bool | int | float | complex | Quantity

Examples

>>> import brainstate
>>> import braintrace
>>>
>>> brainstate.environ.set(precision=64)
>>> x = brainstate.random.randn(16, 8)
>>> B = brainstate.random.randn(8, 2)
>>> A = brainstate.random.randn(2, 4)
>>> y = braintrace.lora_matmul(x, B, A, alpha=0.5)
>>> print(y.shape)
(16, 4)
Contents