Skip to main content
API REFERENCE/omnibias-jax

omnibias.jax.scan_equivariant

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.scan_equivariant
Module contract & mathematical conventions

Discrete-orbit equivariant scan (JAX twin; theory 02-08).

Exact steering is gaussian-family only. C_L is a discrete orbit, not SO(2).

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

OrientationBank​

omnibias.jax.scan_equivariant.OrientationBankSource ↗
class OrientationBank

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.

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

NameType & defaultContract
anglesdataclass fieldtuple[float, ...]RequiredSee the operation contract and implementation below.
steerable_orderdataclass fieldint | NoneDefault: NoneSee the operation contract and implementation below.

Fields & methods

anglesfield
angles: tuple[float, ...]

Declared tuple[float, ...] field on OrientationBank. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
steerable_orderfield
steerable_order: int | None = None

Declared int | None field on OrientationBank. Default expression: None.

View source ↗
Read the implementation
class OrientationBank:
    angles: tuple[float, ...]
    steerable_order: int | None = None
FUNCTION

equivariant_scan_apply​

omnibias.jax.equivariant_scan_applySource ↗
equivariant_scan_apply(x: Array, angles: Sequence[float], offsets: BankSpec, *, base: str='gaussian', template: str='grad') -> Array

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
anglespositional-or-keywordSequence[float]RequiredSee the operation contract and implementation below.
offsetspositional-or-keywordBankSpecRequiredSee the operation contract and implementation below.
basekeyword-onlystrDefault: 'gaussian'See the operation contract and implementation below.
templatekeyword-onlystrDefault: 'grad'See the operation contract and implementation below.
ReturnsArray
Regression example · equivariant scan parity

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_equivariant_scan_parity() -> None:
    pytest.importorskip("jax")
    import jax
    import jax.numpy as jnp
    from omnibias.jax.scan_equivariant import equivariant_scan_apply
    from omnibias.jax.scan_equivariant import steerable_basis as steerable_jax

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    assert steerable_basis(1, 2, base="tanh") is None
    assert steerable_jax(1, 2, base="tanh") is None
    angles = (0.0, math.pi / 2)
    offsets = BankSpec.uniform(-0.4, 0.4, 5)
    net = EquivariantScan(
        2,
        OrientationBank(angles),
        offsets,
        base="gaussian",
        dtype=torch.float64,
    )
    x = torch.tensor([[0.2, -0.15], [0.05, 0.3]], dtype=torch.float64)
    y_t = net(x)
    y_j = equivariant_scan_apply(
        jnp.asarray(x.numpy()), angles, offsets, base="gaussian"
    )
    assert np.asarray(y_t.detach().cpu()) == pytest.approx(
        np.asarray(y_j), rel=1e-12, abs=1e-12
    )
python -m pytest tests/test_equivariant_parity.py -k test_equivariant_scan_parity -q
Read the implementation
def equivariant_scan_apply(
    x: Array,
    angles: Sequence[float],
    offsets: BankSpec,
    *,
    base: str = "gaussian",
    template: str = "grad",
) -> Array:
    act, off, scales, tmpl, _taps = init_bias_scan(1, offsets, template=template, base=base)
    outs = []
    for angle in angles:
        c, s = math.cos(float(angle)), math.sin(float(angle))
        w = jnp.zeros((x.shape[-1],), dtype=x.dtype)
        w = w.at[0].set(c).at[1].set(s)
        z = (x * w).sum(axis=-1, keepdims=True)
        outs.append(scan_response(z, off, scales, tmpl, act))
    return jnp.stack(outs, axis=-1)
Related APIsomnibias.jax.init_bias_scan↗omnibias.jax.scan_response↗
Other public import paths
  • omnibias.jax.equivariant_scan_apply
  • omnibias.jax.scan_equivariant.equivariant_scan_apply
FUNCTION

steerable_basis​

omnibias.jax.steerable_basisSource ↗
steerable_basis(order: int, dim: int, *, base: str='gaussian') -> None | tuple[int, int]

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

NameType & defaultContract
orderpositional-or-keywordintRequiredSee the operation contract and implementation below.
dimpositional-or-keywordintRequiredSee the operation contract and implementation below.
basekeyword-onlystrDefault: 'gaussian'See the operation contract and implementation below.
ReturnsNone | tuple[int, int]
Regression example · equivariant scan parity

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_equivariant_scan_parity() -> None:
    pytest.importorskip("jax")
    import jax
    import jax.numpy as jnp
    from omnibias.jax.scan_equivariant import equivariant_scan_apply
    from omnibias.jax.scan_equivariant import steerable_basis as steerable_jax

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    assert steerable_basis(1, 2, base="tanh") is None
    assert steerable_jax(1, 2, base="tanh") is None
    angles = (0.0, math.pi / 2)
    offsets = BankSpec.uniform(-0.4, 0.4, 5)
    net = EquivariantScan(
        2,
        OrientationBank(angles),
        offsets,
        base="gaussian",
        dtype=torch.float64,
    )
    x = torch.tensor([[0.2, -0.15], [0.05, 0.3]], dtype=torch.float64)
    y_t = net(x)
    y_j = equivariant_scan_apply(
        jnp.asarray(x.numpy()), angles, offsets, base="gaussian"
    )
    assert np.asarray(y_t.detach().cpu()) == pytest.approx(
        np.asarray(y_j), rel=1e-12, abs=1e-12
    )
python -m pytest tests/test_equivariant_parity.py -k test_equivariant_scan_parity -q
Read the implementation
def steerable_basis(order: int, dim: int, *, base: str = "gaussian") -> None | tuple[int, int]:
    if str(base).lower() != "gaussian" or dim != 2:
        return None
    return (int(order), 2)
Other public import paths
  • omnibias.jax.steerable_basis
  • omnibias.jax.scan_equivariant.steerable_basis
← Package overviewRunnable learning examples →