omnibias.torch.scan_equivariant

0.5.0rc1Release candidatePython >=3.10Apache-2.0

PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.

import omnibias.torch.scan_equivariant

Module contract & mathematical conventions

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

Exact steering is gaussian-family only. C_L is a discrete orbit, not SO(2). steerable_basis returns None for non-gaussian bases (no fallback).

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

EquivariantScan

omnibias.torch.EquivariantScanSource ↗

class EquivariantScan(nn.Module)

Scan a C_L orientation orbit. Not SO(2) / SO(3).

Parameters

Name	Type & default	Contract
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`bank`positional-or-keyword	`OrientationBank`Required	See the operation contract and implementation below.
`offsets`positional-or-keyword	`BankSpec`Required	See the operation contract and implementation below.
`template`keyword-only	`OpName \| str`Default: 'grad'	See the operation contract and implementation below.
`base`keyword-only	`str`Default: 'gaussian'	See the operation contract and implementation below.
`readout`keyword-only	`Literal['orbit', 'max', 'fourier']`Default: 'orbit'	See the operation contract and implementation below.
`dtype`keyword-only	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

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

Fields & methods

__init__method

__init__(self, dim: int, bank: OrientationBank, offsets: BankSpec, *, template: OpName | str='grad', base: str='gaussian', readout: Literal['orbit', 'max', 'fourier']='orbit', dtype: torch.dtype | None=None) -> None

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

Parameters

Name	Type & default	Contract
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`bank`positional-or-keyword	`OrientationBank`Required	See the operation contract and implementation below.
`offsets`positional-or-keyword	`BankSpec`Required	See the operation contract and implementation below.
`template`keyword-only	`OpName \| str`Default: 'grad'	See the operation contract and implementation below.
`base`keyword-only	`str`Default: 'gaussian'	See the operation contract and implementation below.
`readout`keyword-only	`Literal['orbit', 'max', 'fourier']`Default: 'orbit'	See the operation contract and implementation below.
`dtype`keyword-only	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

Returns: None

Validation & failure cases

When dim < 2, raises ValueError('equivariant scan needs dim >= 2').

Method implementation

    def __init__(
        self,
        dim: int,
        bank: OrientationBank,
        offsets: BankSpec,
        *,
        template: OpName | str = "grad",
        base: str = "gaussian",
        readout: Literal["orbit", "max", "fourier"] = "orbit",
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        if dim < 2:
            raise ValueError("equivariant scan needs dim >= 2")
        self.dim = int(dim)
        self.orientation = bank
        self.base = str(base)
        self.readout = readout
        dt = torch.get_default_dtype() if dtype is None else dtype
        self.scan = BiasScan(1, offsets, template=template, base=base, dtype=dt)

View source ↗

forwardmethod

forward(self, x: Tensor) -> Tensor

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

Parameters

Name	Type & default	Contract
`x`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.

Returns: Tensor

Method implementation

    def forward(self, x: Tensor) -> Tensor:
        responses = []
        for angle in self.orientation.angles:
            c, s = math.cos(angle), math.sin(angle)
            w = torch.zeros(x.shape[-1], dtype=x.dtype, device=x.device)
            w[0] = c
            w[1] = s
            z = (x * w).sum(dim=-1, keepdim=True)
            responses.append(self.scan(z))
        stacked = torch.stack(responses, dim=-1)
        if self.readout == "max":
            return stacked.amax(dim=-1)
        if self.readout == "fourier":
            return stacked.mean(dim=-1)
        return stacked

View source ↗

dimattribute

dim = int(dim)

Public instance attribute assigned by __init__. The source assignment is dim = int(dim); it describes instance state, not a literal constructor default.

View source ↗

orientationattribute

orientation: OrientationBank = bank

Public instance attribute assigned by __init__. The source assignment is orientation = bank; it describes instance state, not a literal constructor default.

View source ↗

baseattribute

base = str(base)

Public instance attribute assigned by __init__. The source assignment is base = str(base); it describes instance state, not a literal constructor default.

View source ↗

readoutattribute

readout: Literal['orbit', 'max', 'fourier'] = readout

Public instance attribute assigned by __init__. The source assignment is readout = readout; it describes instance state, not a literal constructor default.

View source ↗

scanattribute

scan = BiasScan(1, offsets, template=template, base=base, dtype=dt)

Public instance attribute assigned by __init__. The source assignment is scan = BiasScan(1, offsets, template=template, base=base, dtype=dt); it describes instance state, not a literal constructor default.

View source ↗

Read the implementation

class EquivariantScan(nn.Module):
    """Scan a ``C_L`` orientation orbit. Not SO(2) / SO(3)."""

    def __init__(
        self,
        dim: int,
        bank: OrientationBank,
        offsets: BankSpec,
        *,
        template: OpName | str = "grad",
        base: str = "gaussian",
        readout: Literal["orbit", "max", "fourier"] = "orbit",
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        if dim < 2:
            raise ValueError("equivariant scan needs dim >= 2")
        self.dim = int(dim)
        self.orientation = bank
        self.base = str(base)
        self.readout = readout
        dt = torch.get_default_dtype() if dtype is None else dtype
        self.scan = BiasScan(1, offsets, template=template, base=base, dtype=dt)

    def forward(self, x: Tensor) -> Tensor:
        responses = []
        for angle in self.orientation.angles:
            c, s = math.cos(angle), math.sin(angle)
            w = torch.zeros(x.shape[-1], dtype=x.dtype, device=x.device)
            w[0] = c
            w[1] = s
            z = (x * w).sum(dim=-1, keepdim=True)
            responses.append(self.scan(z))
        stacked = torch.stack(responses, dim=-1)
        if self.readout == "max":
            return stacked.amax(dim=-1)
        if self.readout == "fourier":
            return stacked.mean(dim=-1)
        return stacked

Other public import paths

omnibias.torch.EquivariantScan
omnibias.torch.scan_equivariant.EquivariantScan

CLASS

OrientationBank

omnibias.torch.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

Name	Type & default	Contract
`angles`dataclass field	`tuple[float, ...]`Required	See the operation contract and implementation below.
`steerable_order`dataclass field	`int \| None`Default: None	See the operation contract and implementation below.

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

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

CLASS

SteerableBasis

omnibias.torch.scan_equivariant.SteerableBasisSource ↗

class SteerableBasis

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

Name	Type & default	Contract
`order`dataclass field	`int`Required	See the operation contract and implementation below.
`dim`dataclass field	`int`Required	See the operation contract and implementation below.
`matrices`dataclass field	`tuple[tuple[tuple[float, ...], ...], ...]`Required	See the operation contract and implementation below.

Fields & methods

orderfield

order: int

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

View source ↗

dimfield

dim: int

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

View source ↗

matricesfield

matrices: tuple[tuple[tuple[float, ...], ...], ...]

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

View source ↗

Read the implementation

class SteerableBasis:
    order: int
    dim: int
    matrices: tuple[tuple[tuple[float, ...], ...], ...]

FUNCTION

steerable_basis

omnibias.torch.steerable_basisSource ↗

steerable_basis(order: int, dim: int, *, base: str='gaussian') -> SteerableBasis | None

Exact steering coefficients. None for non-gaussian (no approximation).

Parameters

Name	Type & default	Contract
`order`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`base`keyword-only	`str`Default: 'gaussian'	See the operation contract and implementation below.

ReturnsSteerableBasis | None

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") -> SteerableBasis | None:
    """Exact steering coefficients. ``None`` for non-gaussian (no approximation)."""
    if str(base).lower() != "gaussian":
        return None
    if dim != 2:
        return None
    n = int(order)
    # Rotation of 2-D homogeneous harmonics of degree n: (n+1) cartesian
    # monomials mixed by an SO(2) representation. Stored as R(theta) later.
    return SteerableBasis(order=n, dim=2, matrices=())

Other public import paths

omnibias.torch.steerable_basis
omnibias.torch.scan_equivariant.steerable_basis

← Package overview Runnable learning examples →

omnibias.torch.scan_equivariant

EquivariantScan​

Parameters

Fields & methods

Parameters

Parameters

OrientationBank​

Parameters

Fields & methods

SteerableBasis​

Parameters

Fields & methods

steerable_basis​

Parameters

EquivariantScan

OrientationBank

SteerableBasis

steerable_basis