API REFERENCE/omnibias-jax

omnibias.jax.architectures.jetkan

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.jetkan

Module contract & mathematical conventions

Jet-KAN JAX twin (theory 02-03). Functional, algebraically corresponding to torch.

No GrowableOMBU on JAX: refinement is preallocated zero-weight packs. Kolmogorov-Arnold does not justify the architecture.

Numerical scope: shared coefficients and formulas do not establish universal floating-point identity across backends, devices, compiler transforms, or reduction orders. Compare the chosen path with dtype-appropriate tolerances; even float64 is not a general bitwise-equality guarantee. Algebraic initialization identities likewise need a numerical check for the selected coefficients and reduction order.

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

JetKANConfig

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

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

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

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.

View source ↗

packs_per_edgefield

packs_per_edge: int = 3

Declared int field on JetKANConfig. Default expression: 3.

View source ↗

extra_packsfield

extra_packs: int = 2

Declared int field on JetKANConfig. Default expression: 2.

View source ↗

ordersfield

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

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

View source ↗

basefield

base: str = 'tanh'

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

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 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').

Method implementation

    def __post_init__(self) -> None:
        if len(self.widths) < 2:
            raise ValueError("widths must include at least input and output")
        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")

View source ↗

Read the implementation

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

    def __post_init__(self) -> None:
        if len(self.widths) < 2:
            raise ValueError("widths must include at least input and output")
        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")

Other public import paths

• omnibias.jax.architectures.JetKANConfig
• omnibias.jax.architectures.jetkan.JetKANConfig

CLASS

JetKANParams

omnibias.jax.architectures.JetKANParamsSource ↗

class JetKANParams

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.

Per-layer (weights, means, log_scales) with trailing pack axis.

Parameters

Name	Type & default	Contract
weightsdataclass field	tuple[Array, ...]Required	See the operation contract and implementation below.
meansdataclass field	tuple[Array, ...]Required	See the operation contract and implementation below.
log_scalesdataclass field	tuple[Array, ...]Required	See the operation contract and implementation below.
active_gdataclass field	tuple[int, ...]Required	See the operation contract and implementation below.

Fields & methods

weightsfield

weights: tuple[Array, ...]

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

View source ↗

meansfield

means: tuple[Array, ...]

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

View source ↗

log_scalesfield

log_scales: tuple[Array, ...]

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

View source ↗

active_gfield

active_g: tuple[int, ...]

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

View source ↗

Read the implementation

class JetKANParams:
    """Per-layer ``(weights, means, log_scales)`` with trailing pack axis."""

    weights: tuple[Array, ...]
    means: tuple[Array, ...]
    log_scales: tuple[Array, ...]
    active_g: tuple[int, ...]

Other public import paths

• omnibias.jax.architectures.JetKANParams
• omnibias.jax.architectures.jetkan.JetKANParams

FUNCTION

init_jet_kan

omnibias.jax.architectures.init_jet_kanSource ↗

init_jet_kan(config: JetKANConfig) -> tuple[JaxActivationSpec, JetKANParams]

Zero-init extra packs; live packs get uniform weight 1/G.

Parameters

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

Returnstuple[JaxActivationSpec, JetKANParams]

Read the implementation

def init_jet_kan(config: JetKANConfig) -> tuple[JaxActivationSpec, JetKANParams]:
    """Zero-init extra packs; live packs get uniform weight ``1/G``."""
    act = get_activation(config.base)
    g = _capacity(config)
    live = int(config.packs_per_edge)
    weights: list[Array] = []
    means: list[Array] = []
    log_scales: list[Array] = []
    active: list[int] = []
    for n_in, n_out in zip(config.widths[:-1], config.widths[1:], strict=True):
        w = jnp.zeros((int(n_out), int(n_in), g), dtype=jnp.float64)
        w = w.at[..., :live].set(1.0 / float(max(live, 1)))
        weights.append(w)
        means.append(jnp.zeros((int(n_out), int(n_in), g), dtype=jnp.float64))
        log_scales.append(jnp.zeros((int(n_out), int(n_in), g), dtype=jnp.float64))
        active.append(live)
    params = JetKANParams(
        weights=tuple(weights),
        means=tuple(means),
        log_scales=tuple(log_scales),
        active_g=tuple(active),
    )
    return act, params

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

Other public import paths

• omnibias.jax.architectures.init_jet_kan
• omnibias.jax.architectures.jetkan.init_jet_kan

FUNCTION

jet_kan_apply

omnibias.jax.architectures.jet_kan_applySource ↗

jet_kan_apply(params: JetKANParams, x: Array, *, config: JetKANConfig, base: str | JaxActivationSpec | None=None) -> Array

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
paramspositional-or-keyword	JetKANParamsRequired	See the operation contract and implementation below.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
configkeyword-only	JetKANConfigRequired	See the operation contract and implementation below.
basekeyword-only	str | JaxActivationSpec | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

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

Read the implementation

def jet_kan_apply(
    params: JetKANParams,
    x: Array,
    *,
    config: JetKANConfig,
    base: str | JaxActivationSpec | None = None,
) -> Array:
    act = get_activation(config.base) if base is None else (
        get_activation(base) if isinstance(base, str) else base
    )
    h = x
    for i in range(len(params.weights)):
        h = _layer_apply(
            h,
            params.weights[i],
            params.means[i],
            params.log_scales[i],
            orders=config.orders,
            active_g=params.active_g[i],
            act=act,
        )
    return h

Related APIsomnibias.jax.get_activation↗

Other public import paths

• omnibias.jax.architectures.jet_kan_apply
• omnibias.jax.architectures.jetkan.jet_kan_apply

FUNCTION

jet_kan_from_torch_state

omnibias.jax.architectures.jet_kan_from_torch_stateSource ↗

jet_kan_from_torch_state(config: JetKANConfig, state: Sequence[tuple[Any, Any, Any, int]]) -> JetKANParams

state is [(weights, means, log_scales, active_g), ...].

Parameters

Name	Type & default	Contract
configpositional-or-keyword	JetKANConfigRequired	See the operation contract and implementation below.
statepositional-or-keyword	Sequence[tuple[Any, Any, Any, int]]Required	See the operation contract and implementation below.

ReturnsJetKANParams

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

Read the implementation

def jet_kan_from_torch_state(
    config: JetKANConfig,
    state: Sequence[tuple[Any, Any, Any, int]],
) -> JetKANParams:
    """``state`` is ``[(weights, means, log_scales, active_g), ...]``."""
    del config
    weights = tuple(jnp.asarray(w, dtype=jnp.float64) for w, _m, _s, _a in state)
    means = tuple(jnp.asarray(m, dtype=jnp.float64) for _w, m, _s, _a in state)
    log_scales = tuple(jnp.asarray(s, dtype=jnp.float64) for _w, _m, s, _a in state)
    active = tuple(int(a) for _w, _m, _s, a in state)
    return JetKANParams(weights, means, log_scales, active)

Related APIsomnibias.jax.architectures.JetKANParams↗

Other public import paths

• omnibias.jax.architectures.jet_kan_from_torch_state
• omnibias.jax.architectures.jetkan.jet_kan_from_torch_state

FUNCTION

jet_kan_jet

omnibias.jax.architectures.jet_kan_jetSource ↗

jet_kan_jet(params: JetKANParams, x: Array, order: int, direction: Array, *, config: JetKANConfig) -> Array

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
paramspositional-or-keyword	JetKANParamsRequired	See the operation contract and implementation below.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.
directionpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
configkeyword-only	JetKANConfigRequired	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def jet_kan_jet(
    params: JetKANParams,
    x: Array,
    order: int,
    direction: Array,
    *,
    config: JetKANConfig,
) -> Array:
    act = get_activation(config.base)
    x0 = jnp.asarray(x)
    v = jnp.asarray(direction)
    rows = [x0]
    if order >= 1:
        rows.append(v)
    rows.extend(jnp.zeros_like(x0) for _ in range(max(order - 1, 0)))
    jet = jnp.stack(rows[: order + 1], axis=0)
    for i in range(len(params.weights)):
        jet = _layer_directional_jet(
            jet,
            params.weights[i],
            params.means[i],
            params.log_scales[i],
            orders=config.orders,
            active_g=params.active_g[i],
            act=act,
            order=order,
        )
    return jet

Related APIsomnibias.jax.get_activation↗

Other public import paths

• omnibias.jax.architectures.jet_kan_jet
• omnibias.jax.architectures.jetkan.jet_kan_jet

FUNCTION

jet_kan_jet_mv

omnibias.jax.architectures.jet_kan_jet_mvSource ↗

jet_kan_jet_mv(params: JetKANParams, x: Array, total_order: int, *, config: JetKANConfig) -> Array

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
paramspositional-or-keyword	JetKANParamsRequired	See the operation contract and implementation below.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
total_orderpositional-or-keyword	intRequired	See the operation contract and implementation below.
configkeyword-only	JetKANConfigRequired	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def jet_kan_jet_mv(
    params: JetKANParams,
    x: Array,
    total_order: int,
    *,
    config: JetKANConfig,
) -> Array:
    act = get_activation(config.base)
    x0 = jnp.asarray(x)
    dim = int(x0.shape[0])
    jet = identity_jet(x0, total_order)
    for i in range(len(params.weights)):
        jet = _layer_mixed_jet(
            jet,
            params.weights[i],
            params.means[i],
            params.log_scales[i],
            orders=config.orders,
            active_g=params.active_g[i],
            act=act,
            dim=dim,
            order=total_order,
        )
    return jet

Related APIsomnibias.jax.get_activation↗omnibias.jax.identity_jet↗

Other public import paths

• omnibias.jax.architectures.jet_kan_jet_mv
• omnibias.jax.architectures.jetkan.jet_kan_jet_mv

FUNCTION

jetkan_from_band_plan

omnibias.jax.architectures.jetkan_from_band_planSource ↗

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

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

ReturnsJetKANConfig

Read the implementation

def jetkan_from_band_plan(
    plan: BandPlan,
    widths: tuple[int, ...],
    *,
    base: str = "tanh",
) -> JetKANConfig:
    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,
    )

Related APIsomnibias.jax.architectures.JetKANConfig↗

Other public import paths

• omnibias.jax.architectures.jetkan_from_band_plan
• omnibias.jax.architectures.jetkan.jetkan_from_band_plan

FUNCTION

refine_pack

omnibias.jax.architectures.refine_packSource ↗

refine_pack(params: JetKANParams, config: JetKANConfig) -> JetKANParams

Activate the next zero-weight pack on every layer (forward unchanged).

Parameters

Name	Type & default	Contract
paramspositional-or-keyword	JetKANParamsRequired	See the operation contract and implementation below.
configpositional-or-keyword	JetKANConfigRequired	See the operation contract and implementation below.

ReturnsJetKANParams

Explicit validation & failure cases

• When int(a) >= cap, raises RuntimeError('no unused pack slots remain').

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

Read the implementation

def refine_pack(params: JetKANParams, config: JetKANConfig) -> JetKANParams:
    """Activate the next zero-weight pack on every layer (forward unchanged)."""
    cap = _capacity(config)
    active = []
    for a in params.active_g:
        if int(a) >= cap:
            raise RuntimeError("no unused pack slots remain")
        active.append(int(a) + 1)
    return JetKANParams(params.weights, params.means, params.log_scales, tuple(active))

Related APIsomnibias.jax.architectures.JetKANParams↗

Other public import paths

• omnibias.jax.architectures.refine_pack
• omnibias.jax.architectures.jetkan.refine_pack

← Package overviewRunnable learning examples →