API REFERENCE/omnibias-jax

omnibias.jax.architectures.multiscale

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

Module contract & mathematical conventions

Multi-scale, frequency-aware PINN architectures on the closed-form tower (JAX).

Bit-identical twin of omnibias.torch.architectures.multiscale; see that module for the exposition. The two constructions are:

• AdaptiveActivation -- Jagtap-style sigma(n a z) with a trainable slope a, built as a real ActivationSpec through the backend-neutral omnibias.core.spec.tempered combinator, so the tower (n a)^k sigma^{(k)}(n a z) is exact at every order.
• MscaleMLP -- the MscaleDNN band mixture u(x) = sum_j f_j(alpha_j x), whose jet is the sum of the per-band jets.

Everything here is a frozen dataclass registered as a JAX pytree, so the slopes and band weights are ordinary leaves that jax.grad / jax.jit see. Both backends drive the same Faa di Bruno kernel on the same layer chains, so their derivatives agree to double-precision round-off.

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

AdaptiveActivation

omnibias.jax.architectures.AdaptiveActivationSource ↗

class AdaptiveActivation

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.

Trainable-frequency activation sigma(n a z) as a live ActivationSpec.

JAX twin of omnibias.torch.architectures.multiscale.AdaptiveActivation. a is the single pytree leaf: a scalar for one slope per layer (L-LAAF) or a (hidden,) vector for one per unit (N-LAAF), broadcast over the pre-activation. slope_scale is the fixed amplification n.

Build one with make_adaptive_activation rather than by hand, so the initial *effective* slope n a is the value you asked for.

Parameters

Name	Type & default	Contract
adataclass field	ArrayRequired	See the operation contract and implementation below.
basedataclass field	JaxActivationSpecRequired	See the operation contract and implementation below.
slope_scaledataclass field	floatRequired	See the operation contract and implementation below.

Fields & methods

afield

a: Array

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

View source ↗

basefield

base: JaxActivationSpec

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

View source ↗

slope_scalefield

slope_scale: float

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

View source ↗

scaleproperty

scale(self) -> Array

The effective slope n a at the current parameter value.

Returns: Array

Method implementation

    def scale(self) -> Array:
        """The effective slope ``n a`` at the current parameter value."""
        out: Array = self.slope_scale * self.a
        return out

View source ↗

specproperty

spec(self) -> JaxActivationSpec

Closed-form ActivationSpec of sigma(n a z) at the current a.

Returns: JaxActivationSpec

Method implementation

    def spec(self) -> JaxActivationSpec:
        """Closed-form :class:`ActivationSpec` of ``sigma(n a z)`` at the current ``a``."""
        return tempered(
            self.base,
            self.scale,
            scale="unit",
            name=f"adaptive_{self.base.name}",
            operator_role=(
                f"Adaptive-frequency {self.base.name}: sigma(n a z) with trainable "
                f"slope a; tower (n a)^k sigma^(k)(n a z)."
            ),
        )

View source ↗

Read the implementation

class AdaptiveActivation:
    r"""Trainable-frequency activation ``sigma(n a z)`` as a live :class:`ActivationSpec`.

    JAX twin of :class:`omnibias.torch.architectures.multiscale.AdaptiveActivation`.
    ``a`` is the single pytree leaf: a scalar for one slope per layer (L-LAAF) or a
    ``(hidden,)`` vector for one per unit (N-LAAF), broadcast over the
    pre-activation. ``slope_scale`` is the fixed amplification ``n``.

    Build one with :func:`make_adaptive_activation` rather than by hand, so the
    initial *effective* slope ``n a`` is the value you asked for.
    """

    a: Array
    base: JaxActivationSpec
    slope_scale: float

    @property
    def scale(self) -> Array:
        """The effective slope ``n a`` at the current parameter value."""
        out: Array = self.slope_scale * self.a
        return out

    @property
    def spec(self) -> JaxActivationSpec:
        """Closed-form :class:`ActivationSpec` of ``sigma(n a z)`` at the current ``a``."""
        return tempered(
            self.base,
            self.scale,
            scale="unit",
            name=f"adaptive_{self.base.name}",
            operator_role=(
                f"Adaptive-frequency {self.base.name}: sigma(n a z) with trainable "
                f"slope a; tower (n a)^k sigma^(k)(n a z)."
            ),
        )

Related APIsomnibias.core.tempered↗

Other public import paths

• omnibias.jax.architectures.AdaptiveActivation
• omnibias.jax.architectures.multiscale.AdaptiveActivation

CLASS

AdaptiveJetMLP

omnibias.jax.architectures.AdaptiveJetMLPSource ↗

class AdaptiveJetMLP

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 MLP with a trainable-frequency activation per hidden layer (JAX).

JAX twin of omnibias.torch.architectures.multiscale.AdaptiveJetMLP. The slopes are pytree leaves alongside the weights, so jax.grad trains them together, and every input derivative stays closed form because the slope enters through omnibias.core.spec.tempered.

Parameters

Name	Type & default	Contract
weightsdataclass field	tuple[Array, ...]Required	See the operation contract and implementation below.
biasesdataclass field	tuple[Array, ...]Required	See the operation contract and implementation below.
activationsdataclass field	tuple[AdaptiveActivation, ...]Required	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.

Fields & methods

weightsfield

weights: tuple[Array, ...]

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

View source ↗

biasesfield

biases: tuple[Array, ...]

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

View source ↗

activationsfield

activations: tuple[AdaptiveActivation, ...]

Declared tuple[AdaptiveActivation, ...] field on AdaptiveJetMLP. 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 AdaptiveJetMLP. 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 AdaptiveJetMLP. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

depthproperty

depth(self) -> int

Number of hidden (activated) layers.

Returns: int

Method implementation

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

View source ↗

slopesmethod

slopes(self) -> tuple[Array, ...]

Current effective slope n a of each hidden layer.

Returns: tuple[Array, ...]

Method implementation

    def slopes(self) -> tuple[Array, ...]:
        """Current effective slope ``n a`` of each hidden layer."""
        return tuple(act.scale for act in self.activations)

View source ↗

valuemethod

value(self, x: Array) -> Array

Plain network value u(x), shape (..., out_dim).

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 network value ``u(x)``, shape ``(..., out_dim)``."""
        return _value_from_layers(self._layer_specs(), x)

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 AdaptiveJetMLP:
    r"""Deep MLP with a trainable-frequency activation per hidden layer (JAX).

    JAX twin of :class:`omnibias.torch.architectures.multiscale.AdaptiveJetMLP`.
    The slopes are pytree leaves alongside the weights, so ``jax.grad`` trains
    them together, and every input derivative stays closed form because the slope
    enters through :func:`omnibias.core.spec.tempered`.
    """

    weights: tuple[Array, ...]
    biases: tuple[Array, ...]
    activations: tuple[AdaptiveActivation, ...]
    in_dim: int
    out_dim: int

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

    def _layer_specs(self) -> list[LayerSpec]:
        """``(W, b, spec)`` list with each hidden spec rebuilt at the current slope."""
        n = len(self.weights)
        return [
            (
                self.weights[i],
                self.biases[i],
                None if i == n - 1 else self.activations[i].spec,
            )
            for i in range(n)
        ]

    def _point_hidden_jet(self, xi: Array, order: int) -> Array:
        """Single-point jet through every layer *except* the affine readout.

        Shape ``(M, H)``. Inherits the JetMLP layer-specs path with live slopes.
        """
        return _point_hidden_jet_from_layers(self._layer_specs(), xi, order)

    def _apply_readout_jet(self, hidden_jet: Array) -> Array:
        """Push a (possibly batched) hidden jet through the live affine readout."""
        return _apply_readout_jet_from_layers(self._layer_specs(), hidden_jet)

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

    def _check_fastpath(self, max_order: int) -> None:
        """Reject activations without a closed-form tower up to ``max_order``."""
        _check_layer_fastpaths(self._layer_specs(), max_order)

    def slopes(self) -> tuple[Array, ...]:
        """Current effective slope ``n a`` of each hidden layer."""
        return tuple(act.scale for act in self.activations)

    def value(self, x: Array) -> Array:
        """Plain network value ``u(x)``, shape ``(..., out_dim)``."""
        return _value_from_layers(self._layer_specs(), x)

    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_gradient↗omnibias.jax.jet_hessian↗

Other public import paths

• omnibias.jax.architectures.AdaptiveJetMLP
• omnibias.jax.architectures.multiscale.AdaptiveJetMLP

CLASS

MscaleMLP

omnibias.jax.architectures.MscaleMLPSource ↗

class MscaleMLP

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.

MscaleDNN band mixture u(x) = sum_j f_j(alpha_j x) with exact jets (JAX).

JAX twin of omnibias.torch.architectures.multiscale.MscaleMLP. Each band pre-scales the input by alpha_j, which is the same map as scaling the band's first weight matrix, so every band stays an ordinary layer chain and the mixture's jet is the sum of the per-band jets -- exact and closed form at any order.

Parameters

Name	Type & default	Contract
subnetsdataclass field	tuple[JetMLP, ...]Required	See the operation contract and implementation below.
scalesdataclass field	tuple[float, ...]Required	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.

Fields & methods

subnetsfield

subnets: tuple[JetMLP, ...]

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

View source ↗

scalesfield

scales: tuple[float, ...]

Declared tuple[float, ...] field on MscaleMLP. 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 MscaleMLP. 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 MscaleMLP. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

depthproperty

depth(self) -> int

Number of hidden (activated) layers per band.

Returns: int

Method implementation

    def depth(self) -> int:
        """Number of hidden (activated) layers per band."""
        return self.subnets[0].depth

View source ↗

valuemethod

value(self, x: Array) -> Array

Plain mixture value sum_j f_j(alpha_j x), shape (..., out_dim).

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 mixture value ``sum_j f_j(alpha_j x)``, shape ``(..., out_dim)``."""
        total: Array | None = None
        for layers in self._band_layer_specs():
            h = _value_from_layers(layers, x)
            total = h if total is None else total + h
        assert total is not None
        return total

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 MscaleMLP:
    r"""MscaleDNN band mixture ``u(x) = sum_j f_j(alpha_j x)`` with exact jets (JAX).

    JAX twin of :class:`omnibias.torch.architectures.multiscale.MscaleMLP`. Each
    band pre-scales the input by ``alpha_j``, which is the same map as scaling the
    band's first weight matrix, so every band stays an ordinary layer chain and
    the mixture's jet is the sum of the per-band jets -- exact and closed form at
    any order.
    """

    subnets: tuple[JetMLP, ...]
    scales: tuple[float, ...]
    in_dim: int
    out_dim: int

    @property
    def depth(self) -> int:
        """Number of hidden (activated) layers per band."""
        return self.subnets[0].depth

    def _band_layer_specs(self) -> list[list[LayerSpec]]:
        """One ``(W, b, spec)`` chain per band, with ``W_0`` scaled to ``alpha_j W_0``."""
        groups: list[list[LayerSpec]] = []
        for scale, sub in zip(self.scales, self.subnets, strict=True):
            layers = sub._layer_specs()
            w0, b0, act0 = layers[0]
            groups.append([(w0 * scale, b0, act0), *layers[1:]])
        return groups

    def _point_hidden_jet(self, xi: Array, order: int) -> Array:
        """Stack of per-band hidden jets, shape ``(n_bands, M, H)``."""
        bands: list[Array] = []
        for layers in self._band_layer_specs():
            if len(layers) < 2:
                raise ValueError(
                    "each Mscale band needs a hidden layer plus a readout"
                )
            bands.append(mlp_jet_mv(xi, layers[:-1], order))
        return jnp.stack(bands, axis=0)

    def _apply_readout_jet(self, hidden_jet: Array) -> Array:
        """Apply each band's live readout and sum.

        ``hidden_jet`` is ``(n_bands, M, H)`` or ``(B, n_bands, M, H)``.
        """
        band_layers = self._band_layer_specs()
        total: Array | None = None
        batched = hidden_jet.ndim == 4
        for i, layers in enumerate(band_layers):
            w, b, _spec = layers[-1]
            h_i = hidden_jet[:, i] if batched else hidden_jet[i]
            out_i = jnp.tensordot(h_i, w, axes=([-1], [-1]))
            if b is not None:
                if batched:
                    out_i = out_i.at[:, 0, :].add(b)
                else:
                    out_i = out_i.at[0].add(b)
            total = out_i if total is None else total + out_i
        assert total is not None
        return total

    def _point_jet(self, xi: Array, order: int) -> Array:
        """Single-point jet of the mixture: the sum of the per-band jets."""
        return self._apply_readout_jet(self._point_hidden_jet(xi, order))

    def _check_fastpath(self, max_order: int) -> None:
        """Reject activations without a closed-form tower up to ``max_order``."""
        for sub in self.subnets:
            sub._check_fastpath(max_order)

    def value(self, x: Array) -> Array:
        """Plain mixture value ``sum_j f_j(alpha_j x)``, shape ``(..., out_dim)``."""
        total: Array | None = None
        for layers in self._band_layer_specs():
            h = _value_from_layers(layers, x)
            total = h if total is None else total + h
        assert total is not None
        return total

    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_gradient↗omnibias.jax.jet_hessian↗omnibias.jax.mlp_jet_mv↗

Other public import paths

• omnibias.jax.architectures.MscaleMLP
• omnibias.jax.architectures.multiscale.MscaleMLP

FUNCTION

make_adaptive_activation

omnibias.jax.architectures.make_adaptive_activationSource ↗

make_adaptive_activation(base: str | JaxActivationSpec='tanh', *, slope_scale: float=1.0, width: int | None=None, scale_init: float=1.0, dtype: Any=jnp.float64) -> AdaptiveActivation

Build an AdaptiveActivation whose initial effective slope is scale_init.

Parameters

Name	Type & default	Contract
basepositional-or-keyword	str | JaxActivationSpecDefault: 'tanh'	See the operation contract and implementation below.
slope_scalekeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
widthkeyword-only	int | NoneDefault: None	See the operation contract and implementation below.
scale_initkeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
dtypekeyword-only	AnyDefault: jnp.float64	See the operation contract and implementation below.

ReturnsAdaptiveActivation

Explicit validation & failure cases

• When spec.fastpath is None, raises ValueError(f'AdaptiveActivation requires a base activation with a closed-form derivative kernel; activation {spec.name!r} has none.').
• When slope_scale <= 0.0, raises ValueError(f'slope_scale must be > 0, got {slope_scale}').
• When scale_init <= 0.0, raises ValueError(f'scale_init must be > 0, got {scale_init}').
• When width is not None and width < 1, raises ValueError(f'width must be >= 1 when given, got {width}').

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

Read the implementation

def make_adaptive_activation(
    base: str | JaxActivationSpec = "tanh",
    *,
    slope_scale: float = 1.0,
    width: int | None = None,
    scale_init: float = 1.0,
    dtype: Any = jnp.float64,
) -> AdaptiveActivation:
    """Build an :class:`AdaptiveActivation` whose initial effective slope is ``scale_init``."""
    spec = get_activation(base)
    if spec.fastpath is None:
        raise ValueError(
            f"AdaptiveActivation requires a base activation with a closed-form "
            f"derivative kernel; activation {spec.name!r} has none."
        )
    if slope_scale <= 0.0:
        raise ValueError(f"slope_scale must be > 0, got {slope_scale}")
    if scale_init <= 0.0:
        raise ValueError(f"scale_init must be > 0, got {scale_init}")
    if width is not None and width < 1:
        raise ValueError(f"width must be >= 1 when given, got {width}")
    shape: tuple[int, ...] = () if width is None else (width,)
    a = jnp.full(shape, scale_init / float(slope_scale), dtype=dtype)
    return AdaptiveActivation(a=a, base=spec, slope_scale=float(slope_scale))

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

Other public import paths

• omnibias.jax.architectures.make_adaptive_activation
• omnibias.jax.architectures.multiscale.make_adaptive_activation

FUNCTION

make_adaptive_jet_mlp

omnibias.jax.architectures.make_adaptive_jet_mlpSource ↗

make_adaptive_jet_mlp(in_dim: int, hidden: int, out_dim: int=1, depth: int=3, base: str | JaxActivationSpec='tanh', *, slope_scale: float=1.0, granularity: str='layer', scale_init: float=1.0, seed: int=0, weight_init_scale: float=1.0, dtype: Any=jnp.float64) -> AdaptiveJetMLP

Build a randomly-initialised AdaptiveJetMLP (weights as make_jet_mlp).

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: 3	See the operation contract and implementation below.
basepositional-or-keyword	str | JaxActivationSpecDefault: 'tanh'	See the operation contract and implementation below.
slope_scalekeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
granularitykeyword-only	strDefault: 'layer'	See the operation contract and implementation below.
scale_initkeyword-only	floatDefault: 1.0	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.

ReturnsAdaptiveJetMLP

Explicit validation & failure cases

• When granularity not in ('layer', 'neuron'), raises ValueError(f"granularity must be 'layer' or 'neuron', got {granularity!r}").

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

Read the implementation

def make_adaptive_jet_mlp(
    in_dim: int,
    hidden: int,
    out_dim: int = 1,
    depth: int = 3,
    base: str | JaxActivationSpec = "tanh",
    *,
    slope_scale: float = 1.0,
    granularity: str = "layer",
    scale_init: float = 1.0,
    seed: int = 0,
    weight_init_scale: float = 1.0,
    dtype: Any = jnp.float64,
) -> AdaptiveJetMLP:
    """Build a randomly-initialised :class:`AdaptiveJetMLP` (weights as :func:`make_jet_mlp`)."""
    if granularity not in ("layer", "neuron"):
        raise ValueError(f"granularity must be 'layer' or 'neuron', got {granularity!r}")
    plain = make_jet_mlp(
        in_dim,
        hidden,
        out_dim=out_dim,
        depth=depth,
        activation=base,
        seed=seed,
        weight_init_scale=weight_init_scale,
        dtype=dtype,
    )
    activations = tuple(
        make_adaptive_activation(
            base,
            slope_scale=slope_scale,
            width=hidden if granularity == "neuron" else None,
            scale_init=scale_init,
            dtype=dtype,
        )
        for _ in range(depth)
    )
    return AdaptiveJetMLP(
        weights=plain.weights,
        biases=plain.biases,
        activations=activations,
        in_dim=in_dim,
        out_dim=out_dim,
    )

Related APIsomnibias.jax.architectures.AdaptiveJetMLP↗omnibias.jax.architectures.make_adaptive_activation↗omnibias.jax.architectures.make_jet_mlp↗

Other public import paths

• omnibias.jax.architectures.make_adaptive_jet_mlp
• omnibias.jax.architectures.multiscale.make_adaptive_jet_mlp

FUNCTION

make_mscale_mlp

omnibias.jax.architectures.make_mscale_mlpSource ↗

make_mscale_mlp(in_dim: int, hidden: int, out_dim: int=1, depth: int=3, base: str | JaxActivationSpec='tanh', *, scales: Sequence[float]=(1.0, 2.0, 4.0, 8.0), seed: int=0, weight_init_scale: float=1.0, dtype: Any=jnp.float64) -> MscaleMLP

Build a randomly-initialised MscaleMLP.

hidden is the *total* width, split evenly across the bands (the MscaleDNN convention), so the mixture costs about as much as one MLP of that width. Each band is seeded independently, mirroring the torch builder's fresh draws.

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: 3	See the operation contract and implementation below.
basepositional-or-keyword	str | JaxActivationSpecDefault: 'tanh'	See the operation contract and implementation below.
scaleskeyword-only	Sequence[float]Default: (1.0, 2.0, 4.0, 8.0)	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.

ReturnsMscaleMLP

Explicit validation & failure cases

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

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

Read the implementation

def make_mscale_mlp(
    in_dim: int,
    hidden: int,
    out_dim: int = 1,
    depth: int = 3,
    base: str | JaxActivationSpec = "tanh",
    *,
    scales: Sequence[float] = (1.0, 2.0, 4.0, 8.0),
    seed: int = 0,
    weight_init_scale: float = 1.0,
    dtype: Any = jnp.float64,
) -> MscaleMLP:
    """Build a randomly-initialised :class:`MscaleMLP`.

    ``hidden`` is the *total* width, split evenly across the bands (the MscaleDNN
    convention), so the mixture costs about as much as one MLP of that width. Each
    band is seeded independently, mirroring the torch builder's fresh draws.
    """
    band_scales = _as_band_scales(scales)
    if hidden < 1:
        raise ValueError(f"hidden must be >= 1, got {hidden}")
    band_hidden = max(1, hidden // len(band_scales))
    subnets = tuple(
        make_jet_mlp(
            in_dim,
            band_hidden,
            out_dim=out_dim,
            depth=depth,
            activation=base,
            seed=seed + j,
            weight_init_scale=weight_init_scale,
            dtype=dtype,
        )
        for j in range(len(band_scales))
    )
    return MscaleMLP(
        subnets=subnets, scales=band_scales, in_dim=in_dim, out_dim=out_dim
    )

Related APIsomnibias.jax.architectures.MscaleMLP↗omnibias.jax.architectures.make_jet_mlp↗

Other public import paths

• omnibias.jax.architectures.make_mscale_mlp
• omnibias.jax.architectures.multiscale.make_mscale_mlp

← Package overviewRunnable learning examples →