braintrace.lora_matmul

Contents

braintrace.lora_matmul#

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

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 (ArrayLike) – 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 (ArrayLike) – Low-rank matrix \(B\), shape (in_features, rank).

  • A (ArrayLike) – Low-rank matrix \(A\), shape (rank, out_features).

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

  • bias (ArrayLike or None, optional) – Bias vector, shape (out_features,). Default None.

  • b_fn (callable or None, optional) – 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 or None, optional) – 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 or None, optional) – 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:

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

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