API REFERENCE/omnibias-torch

omnibias.torch.architectures.jetkan

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.torch.architectures.jetkan

Module contract & mathematical conventions

Jet-KAN: univariate multi-pack edges with exact jets (theory 02-03).

Each edge is phi(u) = sum_g c_g sigma^(n_g)(alpha_g u + mu_g). Scales are exp(log_scale) so they stay positive. Exactness is of the **model jet**, not the target; the Kolmogorov-Arnold theorem does **not** justify this architecture (its inner functions may be non-smooth).

Directional jets use ~omnibias.torch.jet.compose_jet; mixed partials use ~omnibias.torch.jet_mv.compose_jet_mv (the same Faà di Bruno kernel as ~omnibias.torch.jet_mv.mlp_jet_mv). Do not reimplement Bell polynomials here.

Refinement in this wave is the honest subset: zero-weight pack birth, and optional ~omnibias.torch.growable.GrowableOperatorMultiBiasUnit order growth on torch. Full residual-driven pack birth/death (03-13) stays designed.

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

JetKAN

omnibias.torch.architectures.JetKANSource ↗

class JetKAN(nn.Module)

Stacked univariate multi-pack edges. Kolmogorov-Arnold is not a proof.

Parameters

Name	Type & default	Contract
configpositional-or-keyword	JetKANConfigRequired	See the operation contract and implementation below.
dtypekeyword-only	torch.dtype | NoneDefault: None	See the operation contract and implementation below.

Executed regression · jetkan from band plan pack count

Recorded coverage confirms this test executed the API body. 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_jetkan_from_band_plan_pack_count() -> None:
    from omnibias.core.spectral_design import design_band_plan
    from omnibias.torch.architectures.jetkan import jetkan_from_band_plan

    plan = design_band_plan("sech", xi_lo=1.0, xi_hi=16.0, channels=3, order=1)
    cfg = jetkan_from_band_plan(plan, (2, 1))
    assert cfg.packs_per_edge == 3
    net = JetKAN(cfg, dtype=torch.float64)
    out = net(torch.zeros(2, 2, dtype=torch.float64))
    assert out.shape == (2, 1)

python -m pytest "packages/omnibias-torch/tests/test_jetkan.py::test_jetkan_from_band_plan_pack_count" -q

Regression example · pack birth is bit identical

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_pack_birth_is_bit_identical() -> None:
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(widths=(2, 2, 1), packs_per_edge=2, extra_packs=2, orders=(0, 1))
    net = JetKAN(cfg, dtype=torch.float64)
    x = torch.randn(5, 2, dtype=torch.float64)
    y0 = net(x).detach().clone()
    net.refine("pack")
    y1 = net(x)
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(y0.reshape(-1).tolist(), y1.reshape(-1).tolist(), strict=True)
    )
    assert worst <= 4.0

python -m pytest packages/omnibias-torch/tests/test_jetkan.py -k test_pack_birth_is_bit_identical -q

Fields & methods

__init__method

__init__(self, config: JetKANConfig, *, dtype: torch.dtype | None=None) -> None

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

Parameters

Name	Type & default	Contract
configpositional-or-keyword	JetKANConfigRequired	See the operation contract and implementation below.
dtypekeyword-only	torch.dtype | NoneDefault: None	See the operation contract and implementation below.

Returns: None

Executed regression · jetkan from band plan pack count

Recorded coverage confirms this test executed the API body. 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_jetkan_from_band_plan_pack_count() -> None:
    from omnibias.core.spectral_design import design_band_plan
    from omnibias.torch.architectures.jetkan import jetkan_from_band_plan

    plan = design_band_plan("sech", xi_lo=1.0, xi_hi=16.0, channels=3, order=1)
    cfg = jetkan_from_band_plan(plan, (2, 1))
    assert cfg.packs_per_edge == 3
    net = JetKAN(cfg, dtype=torch.float64)
    out = net(torch.zeros(2, 2, dtype=torch.float64))
    assert out.shape == (2, 1)

python -m pytest "packages/omnibias-torch/tests/test_jetkan.py::test_jetkan_from_band_plan_pack_count" -q

Method implementation

    def __init__(self, config: JetKANConfig, *, dtype: torch.dtype | None = None) -> None:
        super().__init__()
        self.config = config
        dt = torch.get_default_dtype() if dtype is None else dtype
        act = get_activation(config.base)
        layers: list[_JetKANLayer] = []
        for n_in, n_out in zip(config.widths[:-1], config.widths[1:], strict=True):
            layers.append(
                _JetKANLayer(
                    n_in,
                    n_out,
                    orders=config.orders,
                    n_packs=config.packs_per_edge,
                    extra_packs=config.extra_packs,
                    base=act,
                    growable=config.growable,
                    dtype=dt,
                )
            )
        self.layers = nn.ModuleList(layers)

View source ↗

forwardmethod

forward(self, x: Tensor) -> Tensor

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

Parameters

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

Returns: Tensor

Executed regression · jetkan from band plan pack count

Recorded coverage confirms this test executed the API body. 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_jetkan_from_band_plan_pack_count() -> None:
    from omnibias.core.spectral_design import design_band_plan
    from omnibias.torch.architectures.jetkan import jetkan_from_band_plan

    plan = design_band_plan("sech", xi_lo=1.0, xi_hi=16.0, channels=3, order=1)
    cfg = jetkan_from_band_plan(plan, (2, 1))
    assert cfg.packs_per_edge == 3
    net = JetKAN(cfg, dtype=torch.float64)
    out = net(torch.zeros(2, 2, dtype=torch.float64))
    assert out.shape == (2, 1)

python -m pytest "packages/omnibias-torch/tests/test_jetkan.py::test_jetkan_from_band_plan_pack_count" -q

Method implementation

    def forward(self, x: Tensor) -> Tensor:
        h = x
        for layer in self.layers:
            assert isinstance(layer, _JetKANLayer)
            h = layer(h)
        return h

View source ↗

jetmethod

jet(self, x: Tensor, order: int, direction: Tensor) -> Tensor

Directional Taylor jet of the model along x + t direction, via compose_jet.

Coefficients follow ~omnibias.torch.jet.mlp_jet: jet_to_tower recovers raw directional derivatives.

Parameters

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

Returns: Tensor

Validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').

Executed regression · model jet matches autodiff tanh edge

Recorded coverage confirms this test executed the API body. 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_model_jet_matches_autodiff_tanh_edge() -> None:
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(
        widths=(1, 1), packs_per_edge=1, extra_packs=0, orders=(0,), growable=False
    )
    net = JetKAN(cfg, dtype=torch.float64)
    with torch.no_grad():
        net.layers[0].weights.fill_(1.0)
        net.layers[0].means.zero_()
        net.layers[0].log_scales.zero_()
    from omnibias.torch.jet import jet_to_tower

    x0 = torch.tensor([0.3], dtype=torch.float64)
    v = torch.tensor([1.0], dtype=torch.float64)
    tower = jet_to_tower(net.jet(x0, order=2, direction=v))
    x = torch.tensor([0.3], dtype=torch.float64, requires_grad=True)
    y = net(x.unsqueeze(0)).reshape(())
    d1 = torch.autograd.grad(y, x, create_graph=True)[0]
    d2 = torch.autograd.grad(d1, x, create_graph=True)[0]
    assert _ulp_error(float(tower[0].detach()), float(y.detach())) <= 8.0
    assert _ulp_error(float(tower[1].detach()), float(d1.detach())) <= 16.0
    assert _ulp_error(float(tower[2].detach()), float(d2.detach())) <= 64.0

python -m pytest "packages/omnibias-torch/tests/test_jetkan.py::test_model_jet_matches_autodiff_tanh_edge" -q

Method implementation

    def jet(self, x: Tensor, order: int, direction: Tensor) -> Tensor:
        """Directional Taylor jet of the model along ``x + t direction``, via ``compose_jet``.

    Coefficients follow :func:`~omnibias.torch.jet.mlp_jet`: ``jet_to_tower``
    recovers raw directional derivatives.
    """
        x0 = torch.as_tensor(x)
        v = torch.as_tensor(direction)
        if order < 0:
            raise ValueError(f"order must be >= 0, got {order}")
        rows = [x0]
        if order >= 1:
            rows.append(v)
        rows.extend(torch.zeros_like(x0) for _ in range(max(order - 1, 0)))
        jet = torch.stack(rows[: order + 1], dim=0)
        for layer in self.layers:
            assert isinstance(layer, _JetKANLayer)
            jet = layer.directional_jet(jet, order)
        return jet

View source ↗

jet_mvmethod

jet_mv(self, x: Tensor, total_order: int) -> Tensor

All mixed partials of the model to total_order, via compose_jet_mv.

Parameters

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

Returns: Tensor

Validation & failure cases

• When x0.ndim != 1, raises ValueError(f'jet_mv expects a 1-D point, got shape {tuple(x0.shape)}').

Method implementation

    def jet_mv(self, x: Tensor, total_order: int) -> Tensor:
        """All mixed partials of the model to ``total_order``, via ``compose_jet_mv``."""
        x0 = torch.as_tensor(x)
        if x0.ndim != 1:
            raise ValueError(f"jet_mv expects a 1-D point, got shape {tuple(x0.shape)}")
        dim = int(x0.shape[0])
        jet = identity_jet(x0, total_order)
        for layer in self.layers:
            assert isinstance(layer, _JetKANLayer)
            jet = layer.mixed_jet(jet, dim, total_order)
        return jet

View source ↗

refinemethod

refine(self, kind: str='pack') -> None

Zero-weight pack birth, or GrowableOMBU order growth (torch only).

Parameters

Name	Type & default	Contract
kindpositional-or-keyword	strDefault: 'pack'	See the operation contract and implementation below.

Returns: None

Validation & failure cases

• When layer.active_g >= layer.capacity, raises RuntimeError('no unused pack slots remain').
• When not grew, raises RuntimeError('order growth needs growable=True and free K slots').

Executed regression · pack birth is bit identical

Recorded coverage confirms this test executed the API body. 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_pack_birth_is_bit_identical() -> None:
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(widths=(2, 2, 1), packs_per_edge=2, extra_packs=2, orders=(0, 1))
    net = JetKAN(cfg, dtype=torch.float64)
    x = torch.randn(5, 2, dtype=torch.float64)
    y0 = net(x).detach().clone()
    net.refine("pack")
    y1 = net(x)
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(y0.reshape(-1).tolist(), y1.reshape(-1).tolist(), strict=True)
    )
    assert worst <= 4.0

python -m pytest "packages/omnibias-torch/tests/test_jetkan.py::test_pack_birth_is_bit_identical" -q

Method implementation

    def refine(self, kind: str = "pack") -> None:
        """Zero-weight pack birth, or GrowableOMBU order growth (torch only)."""
        name = str(kind).lower()
        if name == "pack":
            for layer in self.layers:
                assert isinstance(layer, _JetKANLayer)
                if layer.active_g >= layer.capacity:
                    raise RuntimeError("no unused pack slots remain")
                buf = cast(Tensor, layer._active_g)
                buf.copy_(buf + 1)
            return
        if name == "order":
            grew = False
            for layer in self.layers:
                assert isinstance(layer, _JetKANLayer)
                if layer.growable_unit is not None and layer.growable_unit.can_grow:
                    layer.growable_unit.grow(strategy="pair")
                    grew = True
            if not grew:
                raise RuntimeError("order growth needs growable=True and free K slots")
            return
        raise ValueError(f"unknown refine kind {kind!r}; expected 'pack' or 'order'")

View source ↗

configattribute

config: JetKANConfig = config

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

View source ↗

layersattribute

layers = nn.ModuleList(layers)

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

Regression example · model jet matches autodiff tanh edge

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_model_jet_matches_autodiff_tanh_edge() -> None:
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(
        widths=(1, 1), packs_per_edge=1, extra_packs=0, orders=(0,), growable=False
    )
    net = JetKAN(cfg, dtype=torch.float64)
    with torch.no_grad():
        net.layers[0].weights.fill_(1.0)
        net.layers[0].means.zero_()
        net.layers[0].log_scales.zero_()
    from omnibias.torch.jet import jet_to_tower

    x0 = torch.tensor([0.3], dtype=torch.float64)
    v = torch.tensor([1.0], dtype=torch.float64)
    tower = jet_to_tower(net.jet(x0, order=2, direction=v))
    x = torch.tensor([0.3], dtype=torch.float64, requires_grad=True)
    y = net(x.unsqueeze(0)).reshape(())
    d1 = torch.autograd.grad(y, x, create_graph=True)[0]
    d2 = torch.autograd.grad(d1, x, create_graph=True)[0]
    assert _ulp_error(float(tower[0].detach()), float(y.detach())) <= 8.0
    assert _ulp_error(float(tower[1].detach()), float(d1.detach())) <= 16.0
    assert _ulp_error(float(tower[2].detach()), float(d2.detach())) <= 64.0

python -m pytest packages/omnibias-torch/tests/test_jetkan.py -k test_model_jet_matches_autodiff_tanh_edge -q

View source ↗

Read the implementation

class JetKAN(nn.Module):
    """Stacked univariate multi-pack edges. Kolmogorov-Arnold is not a proof."""

    def __init__(self, config: JetKANConfig, *, dtype: torch.dtype | None = None) -> None:
        super().__init__()
        self.config = config
        dt = torch.get_default_dtype() if dtype is None else dtype
        act = get_activation(config.base)
        layers: list[_JetKANLayer] = []
        for n_in, n_out in zip(config.widths[:-1], config.widths[1:], strict=True):
            layers.append(
                _JetKANLayer(
                    n_in,
                    n_out,
                    orders=config.orders,
                    n_packs=config.packs_per_edge,
                    extra_packs=config.extra_packs,
                    base=act,
                    growable=config.growable,
                    dtype=dt,
                )
            )
        self.layers = nn.ModuleList(layers)

    def forward(self, x: Tensor) -> Tensor:
        h = x
        for layer in self.layers:
            assert isinstance(layer, _JetKANLayer)
            h = layer(h)
        return h

    def jet(self, x: Tensor, order: int, direction: Tensor) -> Tensor:
        """Directional Taylor jet of the model along ``x + t direction``, via ``compose_jet``.

    Coefficients follow :func:`~omnibias.torch.jet.mlp_jet`: ``jet_to_tower``
    recovers raw directional derivatives.
    """
        x0 = torch.as_tensor(x)
        v = torch.as_tensor(direction)
        if order < 0:
            raise ValueError(f"order must be >= 0, got {order}")
        rows = [x0]
        if order >= 1:
            rows.append(v)
        rows.extend(torch.zeros_like(x0) for _ in range(max(order - 1, 0)))
        jet = torch.stack(rows[: order + 1], dim=0)
        for layer in self.layers:
            assert isinstance(layer, _JetKANLayer)
            jet = layer.directional_jet(jet, order)
        return jet

    def jet_mv(self, x: Tensor, total_order: int) -> Tensor:
        """All mixed partials of the model to ``total_order``, via ``compose_jet_mv``."""
        x0 = torch.as_tensor(x)
        if x0.ndim != 1:
            raise ValueError(f"jet_mv expects a 1-D point, got shape {tuple(x0.shape)}")
        dim = int(x0.shape[0])
        jet = identity_jet(x0, total_order)
        for layer in self.layers:
            assert isinstance(layer, _JetKANLayer)
            jet = layer.mixed_jet(jet, dim, total_order)
        return jet

    def refine(self, kind: str = "pack") -> None:
        """Zero-weight pack birth, or GrowableOMBU order growth (torch only)."""
        name = str(kind).lower()
        if name == "pack":
            for layer in self.layers:
                assert isinstance(layer, _JetKANLayer)
                if layer.active_g >= layer.capacity:
                    raise RuntimeError("no unused pack slots remain")
                buf = cast(Tensor, layer._active_g)
                buf.copy_(buf + 1)
            return
        if name == "order":
            grew = False
            for layer in self.layers:
                assert isinstance(layer, _JetKANLayer)
                if layer.growable_unit is not None and layer.growable_unit.can_grow:
                    layer.growable_unit.grow(strategy="pair")
                    grew = True
            if not grew:
                raise RuntimeError("order growth needs growable=True and free K slots")
            return
        raise ValueError(f"unknown refine kind {kind!r}; expected 'pack' or 'order'")

Related APIsomnibias.torch.get_activation↗omnibias.torch.identity_jet↗

Other public import paths

• omnibias.torch.architectures.JetKAN
• omnibias.torch.architectures.jetkan.JetKAN

CLASS

JetKANConfig

omnibias.torch.architectures.JetKANConfigSource ↗

class JetKANConfig

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.

Widths include input and output, e.g. (2, 2, 1).

Parameters

Name	Type & default	Contract
widthsdataclass field	tuple[int, ...]Required	See the operation contract and implementation below.
packs_per_edgedataclass field	intDefault: 3	See the operation contract and implementation below.
extra_packsdataclass field	intDefault: 2	See the operation contract and implementation below.
ordersdataclass field	tuple[int, ...]Default: (0, 1, 2)	See the operation contract and implementation below.
basedataclass field	strDefault: 'tanh'	See the operation contract and implementation below.
growabledataclass field	boolDefault: False	See the operation contract and implementation below.

Executed regression · jetkan from band plan pack count

Recorded coverage confirms this test executed the API body. 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_jetkan_from_band_plan_pack_count() -> None:
    from omnibias.core.spectral_design import design_band_plan
    from omnibias.torch.architectures.jetkan import jetkan_from_band_plan

    plan = design_band_plan("sech", xi_lo=1.0, xi_hi=16.0, channels=3, order=1)
    cfg = jetkan_from_band_plan(plan, (2, 1))
    assert cfg.packs_per_edge == 3
    net = JetKAN(cfg, dtype=torch.float64)
    out = net(torch.zeros(2, 2, dtype=torch.float64))
    assert out.shape == (2, 1)

python -m pytest "packages/omnibias-torch/tests/test_jetkan.py::test_jetkan_from_band_plan_pack_count" -q

Regression example · pack birth is bit identical

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_pack_birth_is_bit_identical() -> None:
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(widths=(2, 2, 1), packs_per_edge=2, extra_packs=2, orders=(0, 1))
    net = JetKAN(cfg, dtype=torch.float64)
    x = torch.randn(5, 2, dtype=torch.float64)
    y0 = net(x).detach().clone()
    net.refine("pack")
    y1 = net(x)
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(y0.reshape(-1).tolist(), y1.reshape(-1).tolist(), strict=True)
    )
    assert worst <= 4.0

python -m pytest packages/omnibias-torch/tests/test_jetkan.py -k test_pack_birth_is_bit_identical -q

Fields & methods

widthsfield

widths: tuple[int, ...]

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

Regression example · jetkan torch jax 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_jetkan_torch_jax_parity() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.architectures.jetkan import (
        JetKANConfig as JaxCfg,
    )
    from omnibias.jax.architectures.jetkan import (
        jet_kan_apply,
        jet_kan_from_torch_state,
    )

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(widths=(2, 2, 1), packs_per_edge=2, extra_packs=1, orders=(0, 1))
    net = JetKAN(cfg, dtype=torch.float64)
    x_np = np.array([[0.1, -0.2], [0.3, 0.4], [-0.5, 0.0]], dtype=np.float64)
    t_out = net(torch.as_tensor(x_np)).detach().numpy()
    jcfg = JaxCfg(
        widths=cfg.widths,
        packs_per_edge=cfg.packs_per_edge,
        extra_packs=cfg.extra_packs,
        orders=cfg.orders,
        base=cfg.base,
    )
    params = jet_kan_from_torch_state(jcfg, _torch_state(net))
    j_out = np.asarray(jet_kan_apply(params, jnp.asarray(x_np), config=jcfg))
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(t_out.reshape(-1), j_out.reshape(-1), strict=True)
    )
    assert worst <= 4.0, f"JetKAN parity worst_ulp={worst}"

python -m pytest tests/test_jetkan_parity.py -k test_jetkan_torch_jax_parity -q

View source ↗

packs_per_edgefield

packs_per_edge: int = 3

Declared int field on JetKANConfig. Default expression: 3.

Regression example · jetkan torch jax 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_jetkan_torch_jax_parity() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.architectures.jetkan import (
        JetKANConfig as JaxCfg,
    )
    from omnibias.jax.architectures.jetkan import (
        jet_kan_apply,
        jet_kan_from_torch_state,
    )

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(widths=(2, 2, 1), packs_per_edge=2, extra_packs=1, orders=(0, 1))
    net = JetKAN(cfg, dtype=torch.float64)
    x_np = np.array([[0.1, -0.2], [0.3, 0.4], [-0.5, 0.0]], dtype=np.float64)
    t_out = net(torch.as_tensor(x_np)).detach().numpy()
    jcfg = JaxCfg(
        widths=cfg.widths,
        packs_per_edge=cfg.packs_per_edge,
        extra_packs=cfg.extra_packs,
        orders=cfg.orders,
        base=cfg.base,
    )
    params = jet_kan_from_torch_state(jcfg, _torch_state(net))
    j_out = np.asarray(jet_kan_apply(params, jnp.asarray(x_np), config=jcfg))
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(t_out.reshape(-1), j_out.reshape(-1), strict=True)
    )
    assert worst <= 4.0, f"JetKAN parity worst_ulp={worst}"

python -m pytest tests/test_jetkan_parity.py -k test_jetkan_torch_jax_parity -q

View source ↗

extra_packsfield

extra_packs: int = 2

Declared int field on JetKANConfig. Default expression: 2.

Regression example · jetkan torch jax 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_jetkan_torch_jax_parity() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.architectures.jetkan import (
        JetKANConfig as JaxCfg,
    )
    from omnibias.jax.architectures.jetkan import (
        jet_kan_apply,
        jet_kan_from_torch_state,
    )

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(widths=(2, 2, 1), packs_per_edge=2, extra_packs=1, orders=(0, 1))
    net = JetKAN(cfg, dtype=torch.float64)
    x_np = np.array([[0.1, -0.2], [0.3, 0.4], [-0.5, 0.0]], dtype=np.float64)
    t_out = net(torch.as_tensor(x_np)).detach().numpy()
    jcfg = JaxCfg(
        widths=cfg.widths,
        packs_per_edge=cfg.packs_per_edge,
        extra_packs=cfg.extra_packs,
        orders=cfg.orders,
        base=cfg.base,
    )
    params = jet_kan_from_torch_state(jcfg, _torch_state(net))
    j_out = np.asarray(jet_kan_apply(params, jnp.asarray(x_np), config=jcfg))
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(t_out.reshape(-1), j_out.reshape(-1), strict=True)
    )
    assert worst <= 4.0, f"JetKAN parity worst_ulp={worst}"

python -m pytest tests/test_jetkan_parity.py -k test_jetkan_torch_jax_parity -q

View source ↗

ordersfield

orders: tuple[int, ...] = (0, 1, 2)

Declared tuple[int, ...] field on JetKANConfig. Default expression: (0, 1, 2).

Regression example · jetkan torch jax 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_jetkan_torch_jax_parity() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.architectures.jetkan import (
        JetKANConfig as JaxCfg,
    )
    from omnibias.jax.architectures.jetkan import (
        jet_kan_apply,
        jet_kan_from_torch_state,
    )

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(widths=(2, 2, 1), packs_per_edge=2, extra_packs=1, orders=(0, 1))
    net = JetKAN(cfg, dtype=torch.float64)
    x_np = np.array([[0.1, -0.2], [0.3, 0.4], [-0.5, 0.0]], dtype=np.float64)
    t_out = net(torch.as_tensor(x_np)).detach().numpy()
    jcfg = JaxCfg(
        widths=cfg.widths,
        packs_per_edge=cfg.packs_per_edge,
        extra_packs=cfg.extra_packs,
        orders=cfg.orders,
        base=cfg.base,
    )
    params = jet_kan_from_torch_state(jcfg, _torch_state(net))
    j_out = np.asarray(jet_kan_apply(params, jnp.asarray(x_np), config=jcfg))
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(t_out.reshape(-1), j_out.reshape(-1), strict=True)
    )
    assert worst <= 4.0, f"JetKAN parity worst_ulp={worst}"

python -m pytest tests/test_jetkan_parity.py -k test_jetkan_torch_jax_parity -q

View source ↗

basefield

base: str = 'tanh'

Declared str field on JetKANConfig. Default expression: 'tanh'.

Regression example · jetkan torch jax 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_jetkan_torch_jax_parity() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.jax.architectures.jetkan import (
        JetKANConfig as JaxCfg,
    )
    from omnibias.jax.architectures.jetkan import (
        jet_kan_apply,
        jet_kan_from_torch_state,
    )

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(widths=(2, 2, 1), packs_per_edge=2, extra_packs=1, orders=(0, 1))
    net = JetKAN(cfg, dtype=torch.float64)
    x_np = np.array([[0.1, -0.2], [0.3, 0.4], [-0.5, 0.0]], dtype=np.float64)
    t_out = net(torch.as_tensor(x_np)).detach().numpy()
    jcfg = JaxCfg(
        widths=cfg.widths,
        packs_per_edge=cfg.packs_per_edge,
        extra_packs=cfg.extra_packs,
        orders=cfg.orders,
        base=cfg.base,
    )
    params = jet_kan_from_torch_state(jcfg, _torch_state(net))
    j_out = np.asarray(jet_kan_apply(params, jnp.asarray(x_np), config=jcfg))
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(t_out.reshape(-1), j_out.reshape(-1), strict=True)
    )
    assert worst <= 4.0, f"JetKAN parity worst_ulp={worst}"

python -m pytest tests/test_jetkan_parity.py -k test_jetkan_torch_jax_parity -q

View source ↗

growablefield

growable: bool = False

Declared bool field on JetKANConfig. Default expression: False.

View source ↗

__post_init__method

__post_init__(self) -> None

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

Returns: None

Validation & failure cases

• When len(self.widths) < 2, raises ValueError('widths must include at least input and output').
• When any((int(w) < 1 for w in self.widths)), raises ValueError('each width must be >= 1').
• When int(self.packs_per_edge) < 1, raises ValueError('packs_per_edge must be >= 1').
• When int(self.extra_packs) < 0, raises ValueError('extra_packs must be >= 0').
• When not self.orders, raises ValueError('orders must be non-empty').
• When any((int(n) < 0 for n in self.orders)), raises ValueError('orders must be >= 0').

Executed regression · jetkan from band plan pack count

Recorded coverage confirms this test executed the API body. 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_jetkan_from_band_plan_pack_count() -> None:
    from omnibias.core.spectral_design import design_band_plan
    from omnibias.torch.architectures.jetkan import jetkan_from_band_plan

    plan = design_band_plan("sech", xi_lo=1.0, xi_hi=16.0, channels=3, order=1)
    cfg = jetkan_from_band_plan(plan, (2, 1))
    assert cfg.packs_per_edge == 3
    net = JetKAN(cfg, dtype=torch.float64)
    out = net(torch.zeros(2, 2, dtype=torch.float64))
    assert out.shape == (2, 1)

python -m pytest "packages/omnibias-torch/tests/test_jetkan.py::test_jetkan_from_band_plan_pack_count" -q

Method implementation

    def __post_init__(self) -> None:
        if len(self.widths) < 2:
            raise ValueError("widths must include at least input and output")
        if any(int(w) < 1 for w in self.widths):
            raise ValueError("each width must be >= 1")
        if int(self.packs_per_edge) < 1:
            raise ValueError("packs_per_edge must be >= 1")
        if int(self.extra_packs) < 0:
            raise ValueError("extra_packs must be >= 0")
        if not self.orders:
            raise ValueError("orders must be non-empty")
        if any(int(n) < 0 for n in self.orders):
            raise ValueError("orders must be >= 0")

View source ↗

Read the implementation

class JetKANConfig:
    """Widths include input and output, e.g. ``(2, 2, 1)``."""

    widths: tuple[int, ...]
    packs_per_edge: int = 3
    extra_packs: int = 2
    orders: tuple[int, ...] = (0, 1, 2)
    base: str = "tanh"
    growable: bool = False

    def __post_init__(self) -> None:
        if len(self.widths) < 2:
            raise ValueError("widths must include at least input and output")
        if any(int(w) < 1 for w in self.widths):
            raise ValueError("each width must be >= 1")
        if int(self.packs_per_edge) < 1:
            raise ValueError("packs_per_edge must be >= 1")
        if int(self.extra_packs) < 0:
            raise ValueError("extra_packs must be >= 0")
        if not self.orders:
            raise ValueError("orders must be non-empty")
        if any(int(n) < 0 for n in self.orders):
            raise ValueError("orders must be >= 0")

Other public import paths

• omnibias.torch.architectures.JetKANConfig
• omnibias.torch.architectures.jetkan.JetKANConfig

FUNCTION

edge_functions

omnibias.torch.architectures.edge_functionsSource ↗

edge_functions(model: JetKAN) -> tuple[MultiPackSpec, ...]

Inspect live edge packs (interpretability; means/weights as floats).

Parameters

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

Returnstuple[MultiPackSpec, ...]

Read the implementation

def edge_functions(model: JetKAN) -> tuple[MultiPackSpec, ...]:
    """Inspect live edge packs (interpretability; means/weights as floats)."""
    specs: list[MultiPackSpec] = []
    for layer in model.layers:
        assert isinstance(layer, _JetKANLayer)
        g_live = layer.active_g
        w = layer.weights.detach()
        m = layer.means.detach()
        for q in range(layer.n_out):
            for p in range(layer.n_in):
                packs = []
                for g in range(g_live):
                    n = layer.orders[g] if g < len(layer.orders) else layer.orders[-1]
                    packs.append(
                        PackSpec(
                            order=int(n),
                            mean=float(m[q, p, g]),
                            weight=float(w[q, p, g]),
                        )
                    )
                specs.append(MultiPackSpec(tuple(packs)))
    return tuple(specs)

Related APIsomnibias.core.MultiPackSpec↗omnibias.core.PackSpec↗

Other public import paths

• omnibias.torch.architectures.edge_functions
• omnibias.torch.architectures.jetkan.edge_functions

FUNCTION

jetkan_from_band_plan

omnibias.torch.architectures.jetkan_from_band_planSource ↗

jetkan_from_band_plan(plan: BandPlan, widths: tuple[int, ...], *, base: str='tanh', growable: bool=False) -> JetKANConfig

Optional 01-07 helper: one live pack per band-plan channel.

01-06 wavelet frames stay concept; this only copies the channel count.

Parameters

Name	Type & default	Contract
planpositional-or-keyword	BandPlanRequired	See the operation contract and implementation below.
widthspositional-or-keyword	tuple[int, ...]Required	See the operation contract and implementation below.
basekeyword-only	strDefault: 'tanh'	See the operation contract and implementation below.
growablekeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsJetKANConfig

Executed regression · jetkan from band plan pack count

Recorded coverage confirms this test executed the API body. 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_jetkan_from_band_plan_pack_count() -> None:
    from omnibias.core.spectral_design import design_band_plan
    from omnibias.torch.architectures.jetkan import jetkan_from_band_plan

    plan = design_band_plan("sech", xi_lo=1.0, xi_hi=16.0, channels=3, order=1)
    cfg = jetkan_from_band_plan(plan, (2, 1))
    assert cfg.packs_per_edge == 3
    net = JetKAN(cfg, dtype=torch.float64)
    out = net(torch.zeros(2, 2, dtype=torch.float64))
    assert out.shape == (2, 1)

python -m pytest "packages/omnibias-torch/tests/test_jetkan.py::test_jetkan_from_band_plan_pack_count" -q

Regression example · pack birth is bit identical

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_pack_birth_is_bit_identical() -> None:
    torch.set_default_dtype(torch.float64)
    cfg = JetKANConfig(widths=(2, 2, 1), packs_per_edge=2, extra_packs=2, orders=(0, 1))
    net = JetKAN(cfg, dtype=torch.float64)
    x = torch.randn(5, 2, dtype=torch.float64)
    y0 = net(x).detach().clone()
    net.refine("pack")
    y1 = net(x)
    worst = max(
        _ulp_error(float(a), float(b))
        for a, b in zip(y0.reshape(-1).tolist(), y1.reshape(-1).tolist(), strict=True)
    )
    assert worst <= 4.0

python -m pytest packages/omnibias-torch/tests/test_jetkan.py -k test_pack_birth_is_bit_identical -q

Read the implementation

def jetkan_from_band_plan(
    plan: BandPlan,
    widths: tuple[int, ...],
    *,
    base: str = "tanh",
    growable: bool = False,
) -> JetKANConfig:
    """Optional 01-07 helper: one live pack per band-plan channel.

    01-06 wavelet frames stay concept; this only copies the channel count.
    """
    return JetKANConfig(
        widths=tuple(int(w) for w in widths),
        packs_per_edge=int(plan.n_channels),
        extra_packs=2,
        orders=tuple(int(n) for n in plan.orders),
        base=base,
        growable=growable,
    )

Related APIsomnibias.torch.architectures.JetKANConfig↗

Other public import paths

• omnibias.torch.architectures.jetkan_from_band_plan
• omnibias.torch.architectures.jetkan.jetkan_from_band_plan

← Package overviewRunnable learning examples →