braintrace.element_wise

Contents

braintrace.element_wise#

braintrace.element_wise(weight, *, weight_fn=None)#

ETP-aware element-wise operation.

Applies weight_fn to weight inside the ETP primitive, so the transform is visible to the eligibility-trace compiler. The operation is treated as diagonal in the hidden-state space; the weight participates in trace computation as a per-element trainable parameter.

Parameters:

  • weight (ArrayLike) – Weight parameter (may be a brainunit.Quantity).

  • weight_fn (Callable or None, optional) – Element-wise function applied to the raw weight mantissa inside the primitive. When None (default), the identity is used and the output equals weight exactly. The transform operates on the unitless mantissa; physical units are split off before and recombined after.

Returns:

ArrayLike – weight_fn(weight) (or weight when weight_fn is None), with the same shape as weight.

Examples

>>> import brainstate
>>> import braintrace
>>>
>>> brainstate.environ.set(precision=64)
>>> w = brainstate.random.randn(5)
>>> y = braintrace.element_wise(w)
>>> print(y.shape)
(5,)
>>>
>>> # Apply a non-linearity to the weight
>>> import jax.numpy as jnp
>>> y1 = braintrace.element_wise(w, weight_fn=jnp.tanh)
>>> print(y1.shape)
(5,)
Contents