API REFERENCE/omnibias-jax

omnibias.jax.architectures.attention

0.5.0rc1Release candidatePython >=3.10Apache-2.0

JAX activation derivatives, functional network jets, direct contractions, and parameter-differentiable numerical models.

import omnibias.jax.architectures.attention

Module contract & mathematical conventions

Non-local attention as a closed-form block inside the jet chain (JAX).

Bit-identical twin of omnibias.torch.architectures.attention; see that module for the exposition. The network is

u(x) = W_o\,\Big[\mathrm{softmax}\big(\beta\, q(x) K^\top\big) V + q(x)\Big] + b_o, \qquad q(x) = \mathrm{MLP}(x),

with a trainable memory (K, V) that does not depend on x. The softmax couples all memory slots through a shared denominator, so this is the first genuinely *non-local* block on the substrate -- and omnibias.jax.jet_mv.jet_attention keeps every mixed partial D^alpha u exactly closed form, which is the piece omnibias.hopfield does not supply (it differentiates the log-sum-exp core with respect to the *scores*, not the coordinates).

AttentionJetMLP is a frozen dataclass registered as a JAX pytree, so the encoder weights, the memory, the readout and the inverse temperature are ordinary leaves that jax.grad / jax.jit see.

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →

CLASS

AttentionJetMLP

omnibias.jax.architectures.AttentionJetMLPSource ↗

class AttentionJetMLP

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Deep encoder + one non-local attention block, with exact input derivatives.

JAX twin of omnibias.torch.architectures.attention.AttentionJetMLP. Build one with make_attention_jet_mlp.

Attributes

weights, biases: Encoder chain; every encoder layer is activated by spec. keys, values: Memory of shape (n, d_key) and (n, d_val). readout_weight, readout_bias: Affine readout of shape (out_dim, d_val) and (out_dim,). beta: Inverse temperature (a leaf, so it trains like any other parameter). residual: Whether the query is added to the attention output.

Parameters

Name	Type & default	Contract
weightsdataclass field	tuple[Array, ...]Required	Encoder chain; every encoder layer is activated by spec.
biasesdataclass field	tuple[Array, ...]Required	See the operation contract and implementation below.
keysdataclass field	ArrayRequired	Memory of shape (n, d_key) and (n, d_val).
valuesdataclass field	ArrayRequired	See the operation contract and implementation below.
readout_weightdataclass field	ArrayRequired	Affine readout of shape (out_dim, d_val) and (out_dim,).
readout_biasdataclass field	ArrayRequired	See the operation contract and implementation below.
betadataclass field	ArrayRequired	Inverse temperature (a leaf, so it trains like any other parameter).
specdataclass field	JaxActivationSpecRequired	See the operation contract and implementation below.
in_dimdataclass field	intRequired	See the operation contract and implementation below.
out_dimdataclass field	intRequired	See the operation contract and implementation below.
residualdataclass field	boolDefault: True	Whether the query is added to the attention output.

Fields & methods

weightsfield

weights: tuple[Array, ...]

Encoder chain; every encoder layer is activated by spec.

View source ↗

biasesfield

biases: tuple[Array, ...]

Encoder chain; every encoder layer is activated by spec.

View source ↗

keysfield

keys: Array

Memory of shape (n, d_key) and (n, d_val).

View source ↗

valuesfield

values: Array

Memory of shape (n, d_key) and (n, d_val).

View source ↗

readout_weightfield

readout_weight: Array

Affine readout of shape (out_dim, d_val) and (out_dim,).

View source ↗

readout_biasfield

readout_bias: Array

Affine readout of shape (out_dim, d_val) and (out_dim,).

View source ↗

betafield

beta: Array

Inverse temperature (a leaf, so it trains like any other parameter).

View source ↗

specfield

spec: JaxActivationSpec

Declared JaxActivationSpec field on AttentionJetMLP. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

in_dimfield

in_dim: int

Declared int field on AttentionJetMLP. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

out_dimfield

out_dim: int

Declared int field on AttentionJetMLP. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

residualfield

residual: bool = True

Whether the query is added to the attention output.

View source ↗

depthproperty

depth(self) -> int

Number of hidden (activated) encoder layers.

Returns: int

Method implementation

    def depth(self) -> int:
        """Number of hidden (activated) encoder layers."""
        return len(self.weights)

View source ↗

memoryproperty

memory(self) -> int

Number of memory slots.

Returns: int

Method implementation

    def memory(self) -> int:
        """Number of memory slots."""
        return int(self.keys.shape[0])

View source ↗

value_dimproperty

value_dim(self) -> int

Width of the value slots.

Returns: int

Method implementation

    def value_dim(self) -> int:
        """Width of the value slots."""
        return int(self.values.shape[-1])

View source ↗

querymethod

query(self, x: Array) -> Array

Encoder output q(x), shape (..., hidden).

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

Returns: Array

Method implementation

    def query(self, x: Array) -> Array:
        """Encoder output ``q(x)``, shape ``(..., hidden)``."""
        h = x
        for w, b, spec in self._encoder_specs():
            h = h @ w.T
            if b is not None:
                h = h + b
            assert spec is not None
            h = spec.forward(h)
        out: Array = h
        return out

View source ↗

attention_weightsmethod

attention_weights(self, x: Array) -> Array

Softmax weights over the memory slots, shape (..., memory).

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

Returns: Array

Method implementation

    def attention_weights(self, x: Array) -> Array:
        """Softmax weights over the memory slots, shape ``(..., memory)``."""
        scores = self.query(x) @ (self.keys * self.beta).T
        out: Array = jax.nn.softmax(scores, axis=-1)
        return out

View source ↗

valuemethod

value(self, x: Array) -> Array

Plain forward value u(x), shape (..., out_dim) (no jet needed).

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

Returns: Array

Method implementation

    def value(self, x: Array) -> Array:
        """Plain forward value ``u(x)``, shape ``(..., out_dim)`` (no jet needed)."""
        q = self.query(x)
        h = self.attention_weights(x) @ self.values
        if self.residual:
            h = h + q
        out: Array = h @ self.readout_weight.T + self.readout_bias
        return out

View source ↗

__call__method

__call__(self, x: Array) -> Array

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

Returns: Array

Method implementation

    def __call__(self, x: Array) -> Array:
        return self.value(x)

View source ↗

jetmethod

jet(self, x: Array, order: int) -> Array

Batched multivariate jet, shape (B, M, out_dim).

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returns: Array

Method implementation

    def jet(self, x: Array, order: int) -> Array:
        """Batched multivariate jet, shape ``(B, M, out_dim)``."""
        out: Array = jax.vmap(lambda xi: self._point_jet(xi, order))(x)
        return out

View source ↗

gradientmethod

gradient(self, x: Array) -> Array

Exact input gradient d u / d x_i, shape (B, in_dim, out_dim).

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

Returns: Array

Method implementation

    def gradient(self, x: Array) -> Array:
        """Exact input gradient ``d u / d x_i``, shape ``(B, in_dim, out_dim)``."""
        dim = self.in_dim
        out: Array = jax.vmap(
            lambda xi: jet_gradient(self._point_jet(xi, 1), dim, 1)
        )(x)
        return out

View source ↗

hessianmethod

hessian(self, x: Array) -> Array

Exact input Hessian, shape (B, in_dim, in_dim, out_dim).

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

Returns: Array

Method implementation

    def hessian(self, x: Array) -> Array:
        """Exact input Hessian, shape ``(B, in_dim, in_dim, out_dim)``."""
        dim = self.in_dim
        out: Array = jax.vmap(
            lambda xi: jet_hessian(self._point_jet(xi, 2), dim, 2)
        )(x)
        return out

View source ↗

partialsmethod

partials(self, x: Array, order: int) -> dict[tuple[int, ...], Array]

All raw partials {alpha: D^alpha u(x)} to total order ((B, out_dim)).

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returns: dict[tuple[int, ...], Array]

Method implementation

    def partials(self, x: Array, order: int) -> dict[tuple[int, ...], Array]:
        """All raw partials ``{alpha: D^alpha u(x)}`` to total ``order`` (``(B, out_dim)``)."""
        return _partials_from_jet(self.jet(x, order), self.in_dim, order)

View source ↗

Read the implementation

class AttentionJetMLP:
    r"""Deep encoder + one non-local attention block, with exact input derivatives.

    JAX twin of :class:`omnibias.torch.architectures.attention.AttentionJetMLP`.
    Build one with :func:`make_attention_jet_mlp`.

    Attributes
    ----------
    weights, biases:
        Encoder chain; every encoder layer is activated by ``spec``.
    keys, values:
        Memory of shape ``(n, d_key)`` and ``(n, d_val)``.
    readout_weight, readout_bias:
        Affine readout of shape ``(out_dim, d_val)`` and ``(out_dim,)``.
    beta:
        Inverse temperature (a leaf, so it trains like any other parameter).
    residual:
        Whether the query is added to the attention output.
    """

    weights: tuple[Array, ...]
    biases: tuple[Array, ...]
    keys: Array
    values: Array
    readout_weight: Array
    readout_bias: Array
    beta: Array
    spec: JaxActivationSpec
    in_dim: int
    out_dim: int
    residual: bool = True

    @property
    def depth(self) -> int:
        """Number of hidden (activated) encoder layers."""
        return len(self.weights)

    @property
    def memory(self) -> int:
        """Number of memory slots."""
        return int(self.keys.shape[0])

    @property
    def value_dim(self) -> int:
        """Width of the value slots."""
        return int(self.values.shape[-1])

    def _encoder_specs(self) -> list[LayerSpec]:
        """``(W, b, spec)`` chain of the query encoder ``q(x)`` (all layers activated)."""
        return [
            (w, b, self.spec)
            for w, b in zip(self.weights, self.biases, strict=True)
        ]

    def _check_fastpath(self, max_order: int) -> None:
        """Reject an encoder activation without a closed-form tower of ``max_order``.

        The attention block itself needs no check: ``exp`` and the reciprocal
        tower are closed form at every order by construction.
        """
        _check_layer_fastpaths(self._encoder_specs(), max_order)

    def _point_hidden_jet(self, xi: Array, order: int) -> Array:
        """Single-point pre-readout jet (encoder + attention [+ residual])."""
        q_jet = mlp_jet_mv(xi, self._encoder_specs(), order)
        h_jet = jet_attention(
            q_jet, self.keys, self.values, self.in_dim, order, beta=self.beta
        )
        if self.residual:
            h_jet = h_jet + q_jet
        return h_jet

    def _apply_readout_jet(self, hidden_jet: Array) -> Array:
        """Push a (possibly batched) pre-readout jet through the live affine head."""
        w = self.readout_weight
        b = self.readout_bias
        out = jnp.tensordot(hidden_jet, w, axes=([-1], [-1]))
        if b is not None:
            if hidden_jet.ndim == 2:
                out = out.at[0].add(b)
            else:
                out = out.at[:, 0, :].add(b)
        return out

    def _point_jet(self, xi: Array, order: int) -> Array:
        """Single-point jet of the whole block, shape ``(M, out_dim)``."""
        return self._apply_readout_jet(self._point_hidden_jet(xi, order))

    def query(self, x: Array) -> Array:
        """Encoder output ``q(x)``, shape ``(..., hidden)``."""
        h = x
        for w, b, spec in self._encoder_specs():
            h = h @ w.T
            if b is not None:
                h = h + b
            assert spec is not None
            h = spec.forward(h)
        out: Array = h
        return out

    def attention_weights(self, x: Array) -> Array:
        """Softmax weights over the memory slots, shape ``(..., memory)``."""
        scores = self.query(x) @ (self.keys * self.beta).T
        out: Array = jax.nn.softmax(scores, axis=-1)
        return out

    def value(self, x: Array) -> Array:
        """Plain forward value ``u(x)``, shape ``(..., out_dim)`` (no jet needed)."""
        q = self.query(x)
        h = self.attention_weights(x) @ self.values
        if self.residual:
            h = h + q
        out: Array = h @ self.readout_weight.T + self.readout_bias
        return out

    def __call__(self, x: Array) -> Array:
        return self.value(x)

    def jet(self, x: Array, order: int) -> Array:
        """Batched multivariate jet, shape ``(B, M, out_dim)``."""
        out: Array = jax.vmap(lambda xi: self._point_jet(xi, order))(x)
        return out

    def gradient(self, x: Array) -> Array:
        """Exact input gradient ``d u / d x_i``, shape ``(B, in_dim, out_dim)``."""
        dim = self.in_dim
        out: Array = jax.vmap(
            lambda xi: jet_gradient(self._point_jet(xi, 1), dim, 1)
        )(x)
        return out

    def hessian(self, x: Array) -> Array:
        """Exact input Hessian, shape ``(B, in_dim, in_dim, out_dim)``."""
        dim = self.in_dim
        out: Array = jax.vmap(
            lambda xi: jet_hessian(self._point_jet(xi, 2), dim, 2)
        )(x)
        return out

    def partials(self, x: Array, order: int) -> dict[tuple[int, ...], Array]:
        """All raw partials ``{alpha: D^alpha u(x)}`` to total ``order`` (``(B, out_dim)``)."""
        return _partials_from_jet(self.jet(x, order), self.in_dim, order)

Related APIsomnibias.jax.jet_attention↗omnibias.jax.jet_gradient↗omnibias.jax.jet_hessian↗omnibias.jax.mlp_jet_mv↗

Other public import paths

• omnibias.jax.architectures.AttentionJetMLP
• omnibias.jax.architectures.attention.AttentionJetMLP

FUNCTION

make_attention_jet_mlp

omnibias.jax.architectures.make_attention_jet_mlpSource ↗

make_attention_jet_mlp(in_dim: int, hidden: int, out_dim: int=1, depth: int=2, base: str | JaxActivationSpec='tanh', *, memory: int=16, value_dim: int | None=None, beta: float=1.0, residual: bool=True, seed: int=0, weight_init_scale: float=1.0, dtype: Any=jnp.float64) -> AttentionJetMLP

Build a randomly-initialised AttentionJetMLP.

Weights follow omnibias.jax.architectures.pinn.make_jet_mlp (scaled normals, zero biases); the keys are scaled by 1/sqrt(d_key) so the initial scores are O(1) and the softmax starts near-uniform rather than saturated.

Parameters

Name	Type & default	Contract
in_dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
hiddenpositional-or-keyword	intRequired	See the operation contract and implementation below.
out_dimpositional-or-keyword	intDefault: 1	See the operation contract and implementation below.
depthpositional-or-keyword	intDefault: 2	See the operation contract and implementation below.
basepositional-or-keyword	str | JaxActivationSpecDefault: 'tanh'	See the operation contract and implementation below.
memorykeyword-only	intDefault: 16	See the operation contract and implementation below.
value_dimkeyword-only	int | NoneDefault: None	See the operation contract and implementation below.
betakeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
residualkeyword-only	boolDefault: True	See the operation contract and implementation below.
seedkeyword-only	intDefault: 0	See the operation contract and implementation below.
weight_init_scalekeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
dtypekeyword-only	AnyDefault: jnp.float64	See the operation contract and implementation below.

ReturnsAttentionJetMLP

Explicit validation & failure cases

• When in_dim < 1 or hidden < 1 or out_dim < 1, raises ValueError('in_dim, hidden and out_dim must all be >= 1').
• When depth < 1, raises ValueError(f'depth (number of hidden layers) must be >= 1, got {depth}').
• When memory < 1, raises ValueError(f'memory must be >= 1, got {memory}').
• When beta <= 0.0, raises ValueError(f'beta must be > 0, got {beta}').
• When d_val < 1, raises ValueError(f'value_dim must be >= 1, got {d_val}').
• When residual and d_val != hidden, raises ValueError(f'residual needs value_dim == hidden, got {d_val} != {hidden}; pass residual=False for an asymmetric block').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Read the implementation

def make_attention_jet_mlp(
    in_dim: int,
    hidden: int,
    out_dim: int = 1,
    depth: int = 2,
    base: str | JaxActivationSpec = "tanh",
    *,
    memory: int = 16,
    value_dim: int | None = None,
    beta: float = 1.0,
    residual: bool = True,
    seed: int = 0,
    weight_init_scale: float = 1.0,
    dtype: Any = jnp.float64,
) -> AttentionJetMLP:
    """Build a randomly-initialised :class:`AttentionJetMLP`.

    Weights follow :func:`omnibias.jax.architectures.pinn.make_jet_mlp` (scaled
    normals, zero biases); the keys are scaled by ``1/sqrt(d_key)`` so the initial
    scores are ``O(1)`` and the softmax starts near-uniform rather than saturated.
    """
    if in_dim < 1 or hidden < 1 or out_dim < 1:
        raise ValueError("in_dim, hidden and out_dim must all be >= 1")
    if depth < 1:
        raise ValueError(f"depth (number of hidden layers) must be >= 1, got {depth}")
    if memory < 1:
        raise ValueError(f"memory must be >= 1, got {memory}")
    if beta <= 0.0:
        raise ValueError(f"beta must be > 0, got {beta}")
    d_val = hidden if value_dim is None else int(value_dim)
    if d_val < 1:
        raise ValueError(f"value_dim must be >= 1, got {d_val}")
    if residual and d_val != hidden:
        raise ValueError(
            f"residual needs value_dim == hidden, got {d_val} != {hidden}; "
            "pass residual=False for an asymmetric block"
        )
    spec = get_activation(base)
    key = jax.random.PRNGKey(seed)
    weights: list[Array] = []
    biases: list[Array] = []
    prev = in_dim
    for _ in range(depth):
        key, wk = jax.random.split(key)
        scale = weight_init_scale / math.sqrt(prev)
        weights.append(jax.random.normal(wk, (hidden, prev), dtype=dtype) * scale)
        biases.append(jnp.zeros((hidden,), dtype=dtype))
        prev = hidden
    key, kk, vk, rk = jax.random.split(key, 4)
    keys = jax.random.normal(kk, (memory, hidden), dtype=dtype) / math.sqrt(hidden)
    values = jax.random.normal(vk, (memory, d_val), dtype=dtype)
    readout_weight = jax.random.normal(rk, (out_dim, d_val), dtype=dtype) * (
        weight_init_scale / math.sqrt(d_val)
    )
    return AttentionJetMLP(
        weights=tuple(weights),
        biases=tuple(biases),
        keys=keys,
        values=values,
        readout_weight=readout_weight,
        readout_bias=jnp.zeros((out_dim,), dtype=dtype),
        beta=jnp.asarray(float(beta), dtype=dtype),
        spec=spec,
        in_dim=in_dim,
        out_dim=out_dim,
        residual=bool(residual),
    )

Related APIsomnibias.jax.get_activation↗omnibias.jax.architectures.AttentionJetMLP↗

Other public import paths

• omnibias.jax.architectures.make_attention_jet_mlp
• omnibias.jax.architectures.attention.make_attention_jet_mlp

← Package overviewRunnable learning examples →