API REFERENCE/omnibias-jax

omnibias.jax.architectures.hardbc

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

Module contract & mathematical conventions

Hard-constraint boundary/initial-condition ansatz (JAX) -- twin of torch.

Bit-identical twin of omnibias.torch.architectures.hardbc. The boundary condition is baked into the architecture,

u(x) = g(x) + b(x)\, N(x),

with a fixed **lift** g that matches the boundary data, a **boundary mask** b that vanishes exactly on the constrained set, and a free network N (a omnibias.jax.architectures.JetMLP / ~omnibias.jax.architectures.FourierFeatureMLP / SIREN). The condition then holds for every parameter value, so the boundary term drops out of the loss.

g and b are polynomials whose exact jets come from omnibias.jax.jet_mv.identity_jet; the product b * N is the multivariate Leibniz rule omnibias.jax.jet_mv.jet_multiply. So every D^alpha u is exact and closed form -- no jax.grad for the differential operator. HardConstraintField is a frozen dataclass registered as a JAX pytree (the wrapped network's arrays are the leaves; the fixed mask/lift are static), so it flows through jax.grad / jax.jit for training.

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

AffineFactor

omnibias.jax.architectures.AffineFactorSource ↗

class AffineFactor

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.

A degree-1 boundary factor scale * x[axis] + offset (vanishes at its root).

Parameters

Name	Type & default	Contract
axisdataclass field	intRequired	See the operation contract and implementation below.
scaledataclass field	floatDefault: 1.0	See the operation contract and implementation below.
offsetdataclass field	floatDefault: 0.0	See the operation contract and implementation below.

Executed regression · empty mask and zero scale rejected

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_empty_mask_and_zero_scale_rejected() -> None:
    with pytest.raises(ValueError, match="at least one factor"):
        BoundaryMask(())
    with pytest.raises(ValueError, match="non-zero"):
        AffineFactor(0, 0.0, 1.0)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_empty_mask_and_zero_scale_rejected" -q

Regression example · mask factor axis out of range

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_mask_factor_axis_out_of_range() -> None:
    net = make_jet_mlp(1, 8, 1, depth=2, seed=13)
    with pytest.raises(ValueError, match="out of range"):
        HardConstraintField(net=net, mask=BoundaryMask((AffineFactor(3, 1.0, 0.0),)))

python -m pytest packages/omnibias-jax/tests/test_hardbc.py -k test_mask_factor_axis_out_of_range -q

Fields & methods

axisfield

axis: int

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

View source ↗

scalefield

scale: float = 1.0

Declared float field on AffineFactor. Default expression: 1.0.

View source ↗

offsetfield

offset: float = 0.0

Declared float field on AffineFactor. Default expression: 0.0.

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 self.axis < 0, raises ValueError(f'axis must be >= 0, got {self.axis}').
• When self.scale == 0.0, raises ValueError('scale must be non-zero (a constant factor is not a mask)').

Executed regression · empty mask and zero scale rejected

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_empty_mask_and_zero_scale_rejected() -> None:
    with pytest.raises(ValueError, match="at least one factor"):
        BoundaryMask(())
    with pytest.raises(ValueError, match="non-zero"):
        AffineFactor(0, 0.0, 1.0)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_empty_mask_and_zero_scale_rejected" -q

Method implementation

    def __post_init__(self) -> None:
        if self.axis < 0:
            raise ValueError(f"axis must be >= 0, got {self.axis}")
        if self.scale == 0.0:
            raise ValueError("scale must be non-zero (a constant factor is not a mask)")

View source ↗

Read the implementation

class AffineFactor:
    """A degree-1 boundary factor ``scale * x[axis] + offset`` (vanishes at its root)."""

    axis: int
    scale: float = 1.0
    offset: float = 0.0

    def __post_init__(self) -> None:
        if self.axis < 0:
            raise ValueError(f"axis must be >= 0, got {self.axis}")
        if self.scale == 0.0:
            raise ValueError("scale must be non-zero (a constant factor is not a mask)")

Other public import paths

• omnibias.jax.architectures.AffineFactor
• omnibias.jax.architectures.hardbc.AffineFactor

CLASS

AffineLift

omnibias.jax.architectures.AffineLiftSource ↗

class AffineLift

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.

Affine lift g(x) = W x + c matching the boundary data exactly.

weight is the (out_dim, D) matrix and bias the (out_dim,) vector, both plain Python floats (no trainable state; the boundary data is fixed). A constant lift is W = 0; a 1-D Dirichlet interpolation is a single non-zero slope.

Parameters

Name	Type & default	Contract
weightdataclass field	tuple[tuple[float, ...], ...]Required	See the operation contract and implementation below.
biasdataclass field	tuple[float, ...]Required	See the operation contract and implementation below.

Executed regression · dirichlet interval endpoints exact

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_dirichlet_interval_endpoints_exact() -> None:
    net = make_jet_mlp(1, 16, 1, depth=3, seed=1)
    field = dirichlet_interval(net, 0.0, 2.0, lower_value=1.0, upper_value=-3.0)
    u = field.value(jnp.array([[0.0], [2.0]]))
    assert jnp.allclose(u[:, 0], jnp.array([1.0, -3.0]), atol=1e-12)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_dirichlet_interval_endpoints_exact" -q

Regression example · lift dim mismatch rejected

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_lift_dim_mismatch_rejected() -> None:
    net = make_jet_mlp(2, 8, 1, depth=2, seed=14)
    mask = BoundaryMask((AffineFactor(0, 1.0, 0.0),))
    bad_lift = AffineLift(((0.0,), (0.0,)), (0.0, 0.0))
    with pytest.raises(ValueError, match="out_dim|in_dim"):
        HardConstraintField(net=net, mask=mask, lift=bad_lift)

python -m pytest packages/omnibias-jax/tests/test_hardbc.py -k test_lift_dim_mismatch_rejected -q

Fields & methods

weightfield

weight: tuple[tuple[float, ...], ...]

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

View source ↗

biasfield

bias: tuple[float, ...]

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

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.weight) != len(self.bias), raises ValueError(f'weight has {len(self.weight)} rows but bias has {len(self.bias)}').
• When any((len(row) != width for row in self.weight)), raises ValueError('weight rows must all have the same length').

Executed regression · dirichlet interval endpoints exact

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_dirichlet_interval_endpoints_exact() -> None:
    net = make_jet_mlp(1, 16, 1, depth=3, seed=1)
    field = dirichlet_interval(net, 0.0, 2.0, lower_value=1.0, upper_value=-3.0)
    u = field.value(jnp.array([[0.0], [2.0]]))
    assert jnp.allclose(u[:, 0], jnp.array([1.0, -3.0]), atol=1e-12)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_dirichlet_interval_endpoints_exact" -q

Method implementation

    def __post_init__(self) -> None:
        if len(self.weight) != len(self.bias):
            raise ValueError(
                f"weight has {len(self.weight)} rows but bias has {len(self.bias)}"
            )
        if self.weight:
            width = len(self.weight[0])
            if any(len(row) != width for row in self.weight):
                raise ValueError("weight rows must all have the same length")

View source ↗

out_dimproperty

out_dim(self) -> int

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

Returns: int

Executed regression · dirichlet interval endpoints exact

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_dirichlet_interval_endpoints_exact() -> None:
    net = make_jet_mlp(1, 16, 1, depth=3, seed=1)
    field = dirichlet_interval(net, 0.0, 2.0, lower_value=1.0, upper_value=-3.0)
    u = field.value(jnp.array([[0.0], [2.0]]))
    assert jnp.allclose(u[:, 0], jnp.array([1.0, -3.0]), atol=1e-12)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_dirichlet_interval_endpoints_exact" -q

Method implementation

    def out_dim(self) -> int:
        return len(self.bias)

View source ↗

in_dimproperty

in_dim(self) -> int

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

Returns: int

Executed regression · dirichlet interval endpoints exact

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_dirichlet_interval_endpoints_exact() -> None:
    net = make_jet_mlp(1, 16, 1, depth=3, seed=1)
    field = dirichlet_interval(net, 0.0, 2.0, lower_value=1.0, upper_value=-3.0)
    u = field.value(jnp.array([[0.0], [2.0]]))
    assert jnp.allclose(u[:, 0], jnp.array([1.0, -3.0]), atol=1e-12)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_dirichlet_interval_endpoints_exact" -q

Method implementation

    def in_dim(self) -> int:
        return len(self.weight[0]) if self.weight else 0

View source ↗

valuemethod

value(self, x: Array) -> Array

Lift value g(x), shape (..., out_dim).

Parameters

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

Returns: Array

Executed regression · dirichlet interval endpoints exact

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_dirichlet_interval_endpoints_exact() -> None:
    net = make_jet_mlp(1, 16, 1, depth=3, seed=1)
    field = dirichlet_interval(net, 0.0, 2.0, lower_value=1.0, upper_value=-3.0)
    u = field.value(jnp.array([[0.0], [2.0]]))
    assert jnp.allclose(u[:, 0], jnp.array([1.0, -3.0]), atol=1e-12)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_dirichlet_interval_endpoints_exact" -q

Method implementation

    def value(self, x: Array) -> Array:
        """Lift value ``g(x)``, shape ``(..., out_dim)``."""
        w = jnp.asarray(self.weight, dtype=x.dtype)
        c = jnp.asarray(self.bias, dtype=x.dtype)
        out: Array = x @ w.T + c
        return out

View source ↗

point_jetmethod

point_jet(self, x0: Array, dim: int, order: int) -> Array

Exact jet of g at x0, shape (M, out_dim).

Parameters

Name	Type & default	Contract
x0positional-or-keyword	ArrayRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returns: Array

Executed regression · partials third order match autograd

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_partials_third_order_match_autograd() -> None:
    net = make_jet_mlp(1, 14, 1, depth=2, seed=8)
    field = dirichlet_interval(net, 0.0, 1.0, lower_value=1.0, upper_value=0.0)
    x = jnp.array([[0.37]])
    u = _scalar_value(field)
    d3 = jax.grad(lambda xi: jax.grad(lambda xj: jax.grad(u)(xj)[0])(xi)[0])(x[0])
    assert jnp.allclose(field.partials(x, 3)[(3,)][:, 0], d3, atol=1e-9)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_partials_third_order_match_autograd" -q

Method implementation

    def point_jet(self, x0: Array, dim: int, order: int) -> Array:
        """Exact jet of ``g`` at ``x0``, shape ``(M, out_dim)``."""
        idj = identity_jet(x0, order)  # (M, D)
        w = jnp.asarray(self.weight, dtype=x0.dtype)  # (out, D)
        c = jnp.asarray(self.bias, dtype=x0.dtype)  # (out,)
        g = idj @ w.T  # (M, out)
        m = g.shape[0]
        e0 = (jnp.arange(m) == 0).astype(g.dtype)
        out: Array = g + e0[:, None] * c
        return out

View source ↗

Read the implementation

class AffineLift:
    """Affine lift ``g(x) = W x + c`` matching the boundary data exactly.

    ``weight`` is the ``(out_dim, D)`` matrix and ``bias`` the ``(out_dim,)``
    vector, both plain Python floats (no trainable state; the boundary data is
    fixed). A constant lift is ``W = 0``; a 1-D Dirichlet interpolation is a single
    non-zero slope.
    """

    weight: tuple[tuple[float, ...], ...]
    bias: tuple[float, ...]

    def __post_init__(self) -> None:
        if len(self.weight) != len(self.bias):
            raise ValueError(
                f"weight has {len(self.weight)} rows but bias has {len(self.bias)}"
            )
        if self.weight:
            width = len(self.weight[0])
            if any(len(row) != width for row in self.weight):
                raise ValueError("weight rows must all have the same length")

    @property
    def out_dim(self) -> int:
        return len(self.bias)

    @property
    def in_dim(self) -> int:
        return len(self.weight[0]) if self.weight else 0

    def value(self, x: Array) -> Array:
        """Lift value ``g(x)``, shape ``(..., out_dim)``."""
        w = jnp.asarray(self.weight, dtype=x.dtype)
        c = jnp.asarray(self.bias, dtype=x.dtype)
        out: Array = x @ w.T + c
        return out

    def point_jet(self, x0: Array, dim: int, order: int) -> Array:
        """Exact jet of ``g`` at ``x0``, shape ``(M, out_dim)``."""
        idj = identity_jet(x0, order)  # (M, D)
        w = jnp.asarray(self.weight, dtype=x0.dtype)  # (out, D)
        c = jnp.asarray(self.bias, dtype=x0.dtype)  # (out,)
        g = idj @ w.T  # (M, out)
        m = g.shape[0]
        e0 = (jnp.arange(m) == 0).astype(g.dtype)
        out: Array = g + e0[:, None] * c
        return out

Related APIsomnibias.jax.identity_jet↗

Other public import paths

• omnibias.jax.architectures.AffineLift
• omnibias.jax.architectures.hardbc.AffineLift

CLASS

BoundaryMask

omnibias.jax.architectures.BoundaryMaskSource ↗

class BoundaryMask

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.

Product of affine factors b(x) = \prod_k (s_k x[a_k] + o_k).

Parameters

Name	Type & default	Contract
factorsdataclass field	tuple[AffineFactor, ...]Required	See the operation contract and implementation below.

Executed regression · empty mask and zero scale rejected

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_empty_mask_and_zero_scale_rejected() -> None:
    with pytest.raises(ValueError, match="at least one factor"):
        BoundaryMask(())
    with pytest.raises(ValueError, match="non-zero"):
        AffineFactor(0, 0.0, 1.0)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_empty_mask_and_zero_scale_rejected" -q

Regression example · mask factor axis out of range

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_mask_factor_axis_out_of_range() -> None:
    net = make_jet_mlp(1, 8, 1, depth=2, seed=13)
    with pytest.raises(ValueError, match="out of range"):
        HardConstraintField(net=net, mask=BoundaryMask((AffineFactor(3, 1.0, 0.0),)))

python -m pytest packages/omnibias-jax/tests/test_hardbc.py -k test_mask_factor_axis_out_of_range -q

Fields & methods

factorsfield

factors: tuple[AffineFactor, ...]

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

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 not self.factors, raises ValueError('BoundaryMask needs at least one factor').

Executed regression · empty mask and zero scale rejected

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_empty_mask_and_zero_scale_rejected() -> None:
    with pytest.raises(ValueError, match="at least one factor"):
        BoundaryMask(())
    with pytest.raises(ValueError, match="non-zero"):
        AffineFactor(0, 0.0, 1.0)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_empty_mask_and_zero_scale_rejected" -q

Method implementation

    def __post_init__(self) -> None:
        if not self.factors:
            raise ValueError("BoundaryMask needs at least one factor")

View source ↗

valuemethod

value(self, x: Array) -> Array

Mask value b(x), shape (...,) for input x of shape (..., D).

Parameters

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

Returns: Array

Executed regression · dirichlet homogeneous has no lift

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_dirichlet_homogeneous_has_no_lift() -> None:
    net = make_jet_mlp(1, 8, 1, depth=2, seed=2)
    field = dirichlet_interval(net, -1.0, 1.0)
    assert field.lift is None
    assert jnp.allclose(field.value(jnp.array([[-1.0], [1.0]])), jnp.zeros((2, 1)), atol=1e-13)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_dirichlet_homogeneous_has_no_lift" -q

Method implementation

    def value(self, x: Array) -> Array:
        """Mask value ``b(x)``, shape ``(...,)`` for input ``x`` of shape ``(..., D)``."""
        f0 = self.factors[0]
        out = f0.scale * x[..., f0.axis] + f0.offset
        for f in self.factors[1:]:
            out = out * (f.scale * x[..., f.axis] + f.offset)
        result: Array = out
        return result

View source ↗

point_jetmethod

point_jet(self, x0: Array, dim: int, order: int) -> Array

Exact jet of b at x0, shape (M,) (M multi-indices).

Parameters

Name	Type & default	Contract
x0positional-or-keyword	ArrayRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returns: Array

Executed regression · partials third order match autograd

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_partials_third_order_match_autograd() -> None:
    net = make_jet_mlp(1, 14, 1, depth=2, seed=8)
    field = dirichlet_interval(net, 0.0, 1.0, lower_value=1.0, upper_value=0.0)
    x = jnp.array([[0.37]])
    u = _scalar_value(field)
    d3 = jax.grad(lambda xi: jax.grad(lambda xj: jax.grad(u)(xj)[0])(xi)[0])(x[0])
    assert jnp.allclose(field.partials(x, 3)[(3,)][:, 0], d3, atol=1e-9)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_partials_third_order_match_autograd" -q

Method implementation

    def point_jet(self, x0: Array, dim: int, order: int) -> Array:
        """Exact jet of ``b`` at ``x0``, shape ``(M,)`` (``M`` multi-indices)."""
        idj = identity_jet(x0, order)  # (M, D)
        m = idj.shape[0]
        e0 = (jnp.arange(m) == 0).astype(idj.dtype)
        f0 = self.factors[0]
        acc = f0.scale * idj[:, f0.axis] + f0.offset * e0
        for f in self.factors[1:]:
            fac = f.scale * idj[:, f.axis] + f.offset * e0
            acc = jet_multiply(acc, fac, dim, order)
        result: Array = acc
        return result

View source ↗

Read the implementation

class BoundaryMask:
    r"""Product of affine factors ``b(x) = \prod_k (s_k x[a_k] + o_k)``."""

    factors: tuple[AffineFactor, ...]

    def __post_init__(self) -> None:
        if not self.factors:
            raise ValueError("BoundaryMask needs at least one factor")

    def value(self, x: Array) -> Array:
        """Mask value ``b(x)``, shape ``(...,)`` for input ``x`` of shape ``(..., D)``."""
        f0 = self.factors[0]
        out = f0.scale * x[..., f0.axis] + f0.offset
        for f in self.factors[1:]:
            out = out * (f.scale * x[..., f.axis] + f.offset)
        result: Array = out
        return result

    def point_jet(self, x0: Array, dim: int, order: int) -> Array:
        """Exact jet of ``b`` at ``x0``, shape ``(M,)`` (``M`` multi-indices)."""
        idj = identity_jet(x0, order)  # (M, D)
        m = idj.shape[0]
        e0 = (jnp.arange(m) == 0).astype(idj.dtype)
        f0 = self.factors[0]
        acc = f0.scale * idj[:, f0.axis] + f0.offset * e0
        for f in self.factors[1:]:
            fac = f.scale * idj[:, f.axis] + f.offset * e0
            acc = jet_multiply(acc, fac, dim, order)
        result: Array = acc
        return result

Related APIsomnibias.jax.identity_jet↗omnibias.jax.jet_multiply↗

Other public import paths

• omnibias.jax.architectures.BoundaryMask
• omnibias.jax.architectures.hardbc.BoundaryMask

CLASS

HardConstraintField

omnibias.jax.architectures.HardConstraintFieldSource ↗

class HardConstraintField

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.

Network wrapped so a boundary/initial condition holds exactly: u = g + b N.

net is a omnibias.jax.architectures.JetMLP or ~omnibias.jax.architectures.FourierFeatureMLP. The field exposes the same exact readouts (value, gradient, hessian, value_grad_hessian, jet, partials) but evaluated on u = g + b N rather than net alone, staying closed form via jet-level add/multiply. Use dirichlet_interval, homogeneous_box, initial_value for the common cases.

Parameters

Name	Type & default	Contract
netdataclass field	JetNetRequired	See the operation contract and implementation below.
maskdataclass field	BoundaryMaskRequired	See the operation contract and implementation below.
liftdataclass field	AffineLift | NoneDefault: None	See the operation contract and implementation below.

Executed regression · mask factor axis out of range

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_mask_factor_axis_out_of_range() -> None:
    net = make_jet_mlp(1, 8, 1, depth=2, seed=13)
    with pytest.raises(ValueError, match="out of range"):
        HardConstraintField(net=net, mask=BoundaryMask((AffineFactor(3, 1.0, 0.0),)))

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_mask_factor_axis_out_of_range" -q

Regression example · lift dim mismatch rejected

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_lift_dim_mismatch_rejected() -> None:
    net = make_jet_mlp(2, 8, 1, depth=2, seed=14)
    mask = BoundaryMask((AffineFactor(0, 1.0, 0.0),))
    bad_lift = AffineLift(((0.0,), (0.0,)), (0.0, 0.0))
    with pytest.raises(ValueError, match="out_dim|in_dim"):
        HardConstraintField(net=net, mask=mask, lift=bad_lift)

python -m pytest packages/omnibias-jax/tests/test_hardbc.py -k test_lift_dim_mismatch_rejected -q

Fields & methods

netfield

net: JetNet

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

View source ↗

maskfield

mask: BoundaryMask

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

View source ↗

liftfield

lift: AffineLift | None = None

Declared AffineLift | None field on HardConstraintField. Default expression: None.

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 f.axis >= in_dim, raises ValueError(f'mask factor axis {f.axis} out of range for in_dim={in_dim}').
• When self.lift.out_dim != out_dim, raises ValueError(f'lift out_dim {self.lift.out_dim} != network out_dim {out_dim}').
• When self.lift.in_dim != in_dim, raises ValueError(f'lift in_dim {self.lift.in_dim} != network in_dim {in_dim}').

Executed regression · mask factor axis out of range

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_mask_factor_axis_out_of_range() -> None:
    net = make_jet_mlp(1, 8, 1, depth=2, seed=13)
    with pytest.raises(ValueError, match="out of range"):
        HardConstraintField(net=net, mask=BoundaryMask((AffineFactor(3, 1.0, 0.0),)))

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_mask_factor_axis_out_of_range" -q

Method implementation

    def __post_init__(self) -> None:
        in_dim = int(self.net.in_dim)
        out_dim = int(self.net.out_dim)
        for f in self.mask.factors:
            if f.axis >= in_dim:
                raise ValueError(
                    f"mask factor axis {f.axis} out of range for in_dim={in_dim}"
                )
        if self.lift is not None:
            if self.lift.out_dim != out_dim:
                raise ValueError(
                    f"lift out_dim {self.lift.out_dim} != network out_dim {out_dim}"
                )
            if self.lift.in_dim != in_dim:
                raise ValueError(
                    f"lift in_dim {self.lift.in_dim} != network in_dim {in_dim}"
                )

View source ↗

in_dimproperty

in_dim(self) -> int

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

Returns: int

Executed regression · partials third order match autograd

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_partials_third_order_match_autograd() -> None:
    net = make_jet_mlp(1, 14, 1, depth=2, seed=8)
    field = dirichlet_interval(net, 0.0, 1.0, lower_value=1.0, upper_value=0.0)
    x = jnp.array([[0.37]])
    u = _scalar_value(field)
    d3 = jax.grad(lambda xi: jax.grad(lambda xj: jax.grad(u)(xj)[0])(xi)[0])(x[0])
    assert jnp.allclose(field.partials(x, 3)[(3,)][:, 0], d3, atol=1e-9)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_partials_third_order_match_autograd" -q

Method implementation

    def in_dim(self) -> int:
        return int(self.net.in_dim)

View source ↗

out_dimproperty

out_dim(self) -> int

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

Returns: int

Method implementation

    def out_dim(self) -> int:
        return int(self.net.out_dim)

View source ↗

valuemethod

value(self, x: Array) -> Array

Constrained field value u(x), shape (..., out_dim).

Parameters

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

Returns: Array

Executed regression · dirichlet homogeneous has no lift

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_dirichlet_homogeneous_has_no_lift() -> None:
    net = make_jet_mlp(1, 8, 1, depth=2, seed=2)
    field = dirichlet_interval(net, -1.0, 1.0)
    assert field.lift is None
    assert jnp.allclose(field.value(jnp.array([[-1.0], [1.0]])), jnp.zeros((2, 1)), atol=1e-13)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_dirichlet_homogeneous_has_no_lift" -q

Method implementation

    def value(self, x: Array) -> Array:
        """Constrained field value ``u(x)``, shape ``(..., out_dim)``."""
        u = self.mask.value(x)[..., None] * self.net.value(x)
        if self.lift is not None:
            u = u + self.lift.value(x)
        result: Array = u
        return result

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 of u, 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

Executed regression · partials third order match autograd

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_partials_third_order_match_autograd() -> None:
    net = make_jet_mlp(1, 14, 1, depth=2, seed=8)
    field = dirichlet_interval(net, 0.0, 1.0, lower_value=1.0, upper_value=0.0)
    x = jnp.array([[0.37]])
    u = _scalar_value(field)
    d3 = jax.grad(lambda xi: jax.grad(lambda xj: jax.grad(u)(xj)[0])(xi)[0])(x[0])
    assert jnp.allclose(field.partials(x, 3)[(3,)][:, 0], d3, atol=1e-9)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_partials_third_order_match_autograd" -q

Method implementation

    def jet(self, x: Array, order: int) -> Array:
        """Batched multivariate jet of ``u``, 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

Executed regression · value grad hessian consistent

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_value_grad_hessian_consistent() -> None:
    net = make_jet_mlp(2, 10, 1, depth=2, seed=7)
    field = homogeneous_box(net, [0.0, 0.0], [1.0, 1.0])
    x = jax.random.uniform(jax.random.PRNGKey(0), (5, 2), dtype=jnp.float64)
    v, g, h = field.value_grad_hessian(x)
    assert jnp.allclose(v, field.value(x), atol=1e-12)
    assert jnp.allclose(g, field.gradient(x), atol=1e-12)
    assert jnp.allclose(h, field.hessian(x), atol=1e-12)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_value_grad_hessian_consistent" -q

Method implementation

    def gradient(self, x: Array) -> Array:
        """Exact input gradient ``d u / d x_i``, shape ``(B, in_dim, out_dim)``."""
        d = self.in_dim
        out: Array = jax.vmap(lambda xi: jet_gradient(self._point_jet(xi, 1), d, 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

Executed regression · value grad hessian consistent

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_value_grad_hessian_consistent() -> None:
    net = make_jet_mlp(2, 10, 1, depth=2, seed=7)
    field = homogeneous_box(net, [0.0, 0.0], [1.0, 1.0])
    x = jax.random.uniform(jax.random.PRNGKey(0), (5, 2), dtype=jnp.float64)
    v, g, h = field.value_grad_hessian(x)
    assert jnp.allclose(v, field.value(x), atol=1e-12)
    assert jnp.allclose(g, field.gradient(x), atol=1e-12)
    assert jnp.allclose(h, field.hessian(x), atol=1e-12)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_value_grad_hessian_consistent" -q

Method implementation

    def hessian(self, x: Array) -> Array:
        """Exact input Hessian, shape ``(B, in_dim, in_dim, out_dim)``."""
        d = self.in_dim
        out: Array = jax.vmap(lambda xi: jet_hessian(self._point_jet(xi, 2), d, 2))(x)
        return out

View source ↗

value_grad_hessianmethod

value_grad_hessian(self, x: Array) -> tuple[Array, Array, Array]

One jet -> (value, gradient, Hessian) for 2nd-order PDE residuals.

Parameters

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

Returns: tuple[Array, Array, Array]

Executed regression · value grad hessian consistent

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_value_grad_hessian_consistent() -> None:
    net = make_jet_mlp(2, 10, 1, depth=2, seed=7)
    field = homogeneous_box(net, [0.0, 0.0], [1.0, 1.0])
    x = jax.random.uniform(jax.random.PRNGKey(0), (5, 2), dtype=jnp.float64)
    v, g, h = field.value_grad_hessian(x)
    assert jnp.allclose(v, field.value(x), atol=1e-12)
    assert jnp.allclose(g, field.gradient(x), atol=1e-12)
    assert jnp.allclose(h, field.hessian(x), atol=1e-12)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_value_grad_hessian_consistent" -q

Method implementation

    def value_grad_hessian(self, x: Array) -> tuple[Array, Array, Array]:
        """One jet -> ``(value, gradient, Hessian)`` for 2nd-order PDE residuals."""
        d = self.in_dim

        def f(xi: Array) -> tuple[Array, Array, Array]:
            j = self._point_jet(xi, 2)
            return j[0], jet_gradient(j, d, 2), jet_hessian(j, d, 2)

        res = jax.vmap(f)(x)
        value_b: Array = res[0]
        grad_b: Array = res[1]
        hess_b: Array = res[2]
        return value_b, grad_b, hess_b

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]

Executed regression · partials third order match autograd

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_partials_third_order_match_autograd() -> None:
    net = make_jet_mlp(1, 14, 1, depth=2, seed=8)
    field = dirichlet_interval(net, 0.0, 1.0, lower_value=1.0, upper_value=0.0)
    x = jnp.array([[0.37]])
    u = _scalar_value(field)
    d3 = jax.grad(lambda xi: jax.grad(lambda xj: jax.grad(u)(xj)[0])(xi)[0])(x[0])
    assert jnp.allclose(field.partials(x, 3)[(3,)][:, 0], d3, atol=1e-9)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_partials_third_order_match_autograd" -q

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)``)."""
        jet_b = self.jet(x, order)
        idx = multi_indices(self.in_dim, order)
        return {
            alpha: jet_b[:, i] * multi_index_factorial(alpha)
            for i, alpha in enumerate(idx)
        }

View source ↗

Read the implementation

class HardConstraintField:
    r"""Network wrapped so a boundary/initial condition holds exactly: ``u = g + b N``.

    ``net`` is a :class:`omnibias.jax.architectures.JetMLP` or
    :class:`~omnibias.jax.architectures.FourierFeatureMLP`. The field exposes the
    same exact readouts (:meth:`value`, :meth:`gradient`, :meth:`hessian`,
    :meth:`value_grad_hessian`, :meth:`jet`, :meth:`partials`) but evaluated on
    ``u = g + b N`` rather than ``net`` alone, staying closed form via jet-level
    add/multiply. Use :func:`dirichlet_interval`, :func:`homogeneous_box`,
    :func:`initial_value` for the common cases.
    """

    net: JetNet
    mask: BoundaryMask
    lift: AffineLift | None = None

    def __post_init__(self) -> None:
        in_dim = int(self.net.in_dim)
        out_dim = int(self.net.out_dim)
        for f in self.mask.factors:
            if f.axis >= in_dim:
                raise ValueError(
                    f"mask factor axis {f.axis} out of range for in_dim={in_dim}"
                )
        if self.lift is not None:
            if self.lift.out_dim != out_dim:
                raise ValueError(
                    f"lift out_dim {self.lift.out_dim} != network out_dim {out_dim}"
                )
            if self.lift.in_dim != in_dim:
                raise ValueError(
                    f"lift in_dim {self.lift.in_dim} != network in_dim {in_dim}"
                )

    @property
    def in_dim(self) -> int:
        return int(self.net.in_dim)

    @property
    def out_dim(self) -> int:
        return int(self.net.out_dim)

    def _point_jet(self, xi: Array, order: int) -> Array:
        net_jet = mlp_jet_mv(xi, self.net._layer_specs(), order)  # (M, out)
        mask_jet = self.mask.point_jet(xi, self.in_dim, order)  # (M,)
        u_jet = jet_multiply(mask_jet, net_jet, self.in_dim, order)
        if self.lift is not None:
            u_jet = u_jet + self.lift.point_jet(xi, self.in_dim, order)
        return u_jet

    def value(self, x: Array) -> Array:
        """Constrained field value ``u(x)``, shape ``(..., out_dim)``."""
        u = self.mask.value(x)[..., None] * self.net.value(x)
        if self.lift is not None:
            u = u + self.lift.value(x)
        result: Array = u
        return result

    def __call__(self, x: Array) -> Array:
        return self.value(x)

    def jet(self, x: Array, order: int) -> Array:
        """Batched multivariate jet of ``u``, 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)``."""
        d = self.in_dim
        out: Array = jax.vmap(lambda xi: jet_gradient(self._point_jet(xi, 1), d, 1))(x)
        return out

    def hessian(self, x: Array) -> Array:
        """Exact input Hessian, shape ``(B, in_dim, in_dim, out_dim)``."""
        d = self.in_dim
        out: Array = jax.vmap(lambda xi: jet_hessian(self._point_jet(xi, 2), d, 2))(x)
        return out

    def value_grad_hessian(self, x: Array) -> tuple[Array, Array, Array]:
        """One jet -> ``(value, gradient, Hessian)`` for 2nd-order PDE residuals."""
        d = self.in_dim

        def f(xi: Array) -> tuple[Array, Array, Array]:
            j = self._point_jet(xi, 2)
            return j[0], jet_gradient(j, d, 2), jet_hessian(j, d, 2)

        res = jax.vmap(f)(x)
        value_b: Array = res[0]
        grad_b: Array = res[1]
        hess_b: Array = res[2]
        return value_b, grad_b, hess_b

    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)``)."""
        jet_b = self.jet(x, order)
        idx = multi_indices(self.in_dim, order)
        return {
            alpha: jet_b[:, i] * multi_index_factorial(alpha)
            for i, alpha in enumerate(idx)
        }

Related APIsomnibias.core.multi_index_factorial↗omnibias.core.multi_indices↗omnibias.jax.jet_gradient↗omnibias.jax.jet_hessian↗omnibias.jax.jet_multiply↗omnibias.jax.mlp_jet_mv↗

Other public import paths

• omnibias.jax.architectures.HardConstraintField
• omnibias.jax.architectures.hardbc.HardConstraintField

FUNCTION

dirichlet_interval

omnibias.jax.architectures.dirichlet_intervalSource ↗

dirichlet_interval(net: JetNet, lower: float, upper: float, *, lower_value: float=0.0, upper_value: float=0.0, axis: int=0) -> HardConstraintField

Hard 1-D Dirichlet BC: u(lower) = lower_value, u(upper) = upper_value.

Mask b(x) = (x_axis - lower)(upper - x_axis) plus the linear lift interpolating the endpoint values, so the field matches the prescribed data at both ends exactly. The same boundary profile is applied to every output.

Parameters

Name	Type & default	Contract
netpositional-or-keyword	JetNetRequired	See the operation contract and implementation below.
lowerpositional-or-keyword	floatRequired	See the operation contract and implementation below.
upperpositional-or-keyword	floatRequired	See the operation contract and implementation below.
lower_valuekeyword-only	floatDefault: 0.0	See the operation contract and implementation below.
upper_valuekeyword-only	floatDefault: 0.0	See the operation contract and implementation below.
axiskeyword-only	intDefault: 0	See the operation contract and implementation below.

ReturnsHardConstraintField

Explicit validation & failure cases

• When upper <= lower, raises ValueError(f'upper ({upper}) must exceed lower ({lower})').

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

Executed regression · dirichlet interval rejects bad bounds

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_dirichlet_interval_rejects_bad_bounds() -> None:
    net = make_jet_mlp(1, 8, 1, depth=2, seed=11)
    with pytest.raises(ValueError, match="must exceed"):
        dirichlet_interval(net, 1.0, 0.0)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_dirichlet_interval_rejects_bad_bounds" -q

Regression example · dirichlet homogeneous has no lift

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_dirichlet_homogeneous_has_no_lift() -> None:
    net = make_jet_mlp(1, 8, 1, depth=2, seed=2)
    field = dirichlet_interval(net, -1.0, 1.0)
    assert field.lift is None
    assert jnp.allclose(field.value(jnp.array([[-1.0], [1.0]])), jnp.zeros((2, 1)), atol=1e-13)

python -m pytest packages/omnibias-jax/tests/test_hardbc.py -k test_dirichlet_homogeneous_has_no_lift -q

Read the implementation

def dirichlet_interval(
    net: JetNet,
    lower: float,
    upper: float,
    *,
    lower_value: float = 0.0,
    upper_value: float = 0.0,
    axis: int = 0,
) -> HardConstraintField:
    r"""Hard 1-D Dirichlet BC: ``u(lower) = lower_value``, ``u(upper) = upper_value``.

    Mask ``b(x) = (x_axis - lower)(upper - x_axis)`` plus the linear lift
    interpolating the endpoint values, so the field matches the prescribed data at
    both ends exactly. The same boundary profile is applied to every output.
    """
    if upper <= lower:
        raise ValueError(f"upper ({upper}) must exceed lower ({lower})")
    in_dim = int(net.in_dim)
    out_dim = int(net.out_dim)
    mask = BoundaryMask(
        (AffineFactor(axis, 1.0, -lower), AffineFactor(axis, -1.0, upper))
    )
    lift: AffineLift | None = None
    if lower_value != 0.0 or upper_value != 0.0:
        slope = (upper_value - lower_value) / (upper - lower)
        intercept = lower_value - slope * lower
        row = tuple(slope if j == axis else 0.0 for j in range(in_dim))
        lift = AffineLift((row,) * out_dim, (intercept,) * out_dim)
    return HardConstraintField(net=net, mask=mask, lift=lift)

Related APIsomnibias.jax.architectures.AffineFactor↗omnibias.jax.architectures.AffineLift↗omnibias.jax.architectures.BoundaryMask↗omnibias.jax.architectures.HardConstraintField↗

Other public import paths

• omnibias.jax.architectures.dirichlet_interval
• omnibias.jax.architectures.hardbc.dirichlet_interval

FUNCTION

homogeneous_box

omnibias.jax.architectures.homogeneous_boxSource ↗

homogeneous_box(net: JetNet, lows: Sequence[float], highs: Sequence[float]) -> HardConstraintField

Hard homogeneous Dirichlet BC u = 0 on the boundary of a box.

Mask \prod_i (x_i - lows_i)(highs_i - x_i) over the box \prod_i [lows_i, highs_i]; the lift is zero, so u = b N.

Parameters

Name	Type & default	Contract
netpositional-or-keyword	JetNetRequired	See the operation contract and implementation below.
lowspositional-or-keyword	Sequence[float]Required	See the operation contract and implementation below.
highspositional-or-keyword	Sequence[float]Required	See the operation contract and implementation below.

ReturnsHardConstraintField

Explicit validation & failure cases

• When len(lows) != in_dim or len(highs) != in_dim, raises ValueError(f'lows/highs must have length in_dim={in_dim}, got {len(lows)} and {len(highs)}').
• When highs[i] <= lows[i], raises ValueError(f'highs[{i}] ({highs[i]}) must exceed lows[{i}] ({lows[i]})').

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

Executed regression · homogeneous box rejects wrong length

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_homogeneous_box_rejects_wrong_length() -> None:
    net = make_jet_mlp(2, 8, 1, depth=2, seed=12)
    with pytest.raises(ValueError, match="length in_dim"):
        homogeneous_box(net, [0.0], [1.0])

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_homogeneous_box_rejects_wrong_length" -q

Regression example · multi output box 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_multi_output_box_parity() -> None:
    tnet, jnet = _jetmlp_pair(2, 10, 3, 2, "tanh", seed=4)
    tfield = torch_hardbc.homogeneous_box(tnet, [0.0, 0.0], [1.0, 1.0])
    jfield = jax_hardbc.homogeneous_box(jnet, [0.0, 0.0], [1.0, 1.0])
    xt, xj = _xpair(seed=4, b=4, d=2)
    _assert_field_parity(tfield, jfield, xt, xj)

python -m pytest tests/test_hardbc_parity.py -k test_multi_output_box_parity -q

Read the implementation

def homogeneous_box(
    net: JetNet,
    lows: Sequence[float],
    highs: Sequence[float],
) -> HardConstraintField:
    r"""Hard homogeneous Dirichlet BC ``u = 0`` on the boundary of a box.

    Mask ``\prod_i (x_i - lows_i)(highs_i - x_i)`` over the box
    ``\prod_i [lows_i, highs_i]``; the lift is zero, so ``u = b N``.
    """
    in_dim = int(net.in_dim)
    if len(lows) != in_dim or len(highs) != in_dim:
        raise ValueError(
            f"lows/highs must have length in_dim={in_dim}, "
            f"got {len(lows)} and {len(highs)}"
        )
    factors: list[AffineFactor] = []
    for i in range(in_dim):
        if highs[i] <= lows[i]:
            raise ValueError(
                f"highs[{i}] ({highs[i]}) must exceed lows[{i}] ({lows[i]})"
            )
        factors.append(AffineFactor(i, 1.0, -float(lows[i])))
        factors.append(AffineFactor(i, -1.0, float(highs[i])))
    return HardConstraintField(net=net, mask=BoundaryMask(tuple(factors)), lift=None)

Related APIsomnibias.jax.architectures.AffineFactor↗omnibias.jax.architectures.BoundaryMask↗omnibias.jax.architectures.HardConstraintField↗

Other public import paths

• omnibias.jax.architectures.homogeneous_box
• omnibias.jax.architectures.hardbc.homogeneous_box

FUNCTION

initial_value

omnibias.jax.architectures.initial_valueSource ↗

initial_value(net: JetNet, *, t_axis: int=-1, t0: float=0.0, value: float=0.0) -> HardConstraintField

Hard (constant) initial condition u(x, t0) = value.

Mask t - t0 (one-sided: zero only on the initial slice) plus a constant lift. For a non-constant initial profile, build HardConstraintField directly with a custom AffineLift / BoundaryMask.

Parameters

Name	Type & default	Contract
netpositional-or-keyword	JetNetRequired	See the operation contract and implementation below.
t_axiskeyword-only	intDefault: -1	See the operation contract and implementation below.
t0keyword-only	floatDefault: 0.0	See the operation contract and implementation below.
valuekeyword-only	floatDefault: 0.0	See the operation contract and implementation below.

ReturnsHardConstraintField

Explicit validation & failure cases

• When not 0 <= ax < in_dim, raises ValueError(f't_axis {t_axis} out of range for in_dim={in_dim}').

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

Executed regression · initial value exact on slice

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_initial_value_exact_on_slice() -> None:
    net = make_jet_mlp(2, 10, 1, depth=2, seed=6)
    field = initial_value(net, t_axis=1, t0=0.0, value=2.0)
    ic = jnp.array([[0.1, 0.0], [0.9, 0.0], [0.5, 0.0]])
    assert jnp.allclose(field.value(ic), 2.0 * jnp.ones((3, 1)), atol=1e-13)
    assert abs(float(field.value(jnp.array([[0.5, 0.7]]))[0, 0]) - 2.0) > 0.0

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_initial_value_exact_on_slice" -q

Regression example · initial value 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_initial_value_parity() -> None:
    tnet, jnet = _jetmlp_pair(2, 10, 1, 2, "sigmoid", seed=3)
    tfield = torch_hardbc.initial_value(tnet, t_axis=1, t0=0.0, value=2.0)
    jfield = jax_hardbc.initial_value(jnet, t_axis=1, t0=0.0, value=2.0)
    xt, xj = _xpair(seed=3, b=5, d=2)
    _assert_field_parity(tfield, jfield, xt, xj)

python -m pytest tests/test_hardbc_parity.py -k test_initial_value_parity -q

Read the implementation

def initial_value(
    net: JetNet,
    *,
    t_axis: int = -1,
    t0: float = 0.0,
    value: float = 0.0,
) -> HardConstraintField:
    r"""Hard (constant) initial condition ``u(x, t0) = value``.

    Mask ``t - t0`` (one-sided: zero only on the initial slice) plus a constant
    lift. For a non-constant initial profile, build :class:`HardConstraintField`
    directly with a custom :class:`AffineLift` / :class:`BoundaryMask`.
    """
    in_dim = int(net.in_dim)
    out_dim = int(net.out_dim)
    ax = t_axis if t_axis >= 0 else in_dim + t_axis
    if not (0 <= ax < in_dim):
        raise ValueError(f"t_axis {t_axis} out of range for in_dim={in_dim}")
    mask = BoundaryMask((AffineFactor(ax, 1.0, -t0),))
    lift: AffineLift | None = None
    if value != 0.0:
        zeros = (0.0,) * in_dim
        lift = AffineLift((zeros,) * out_dim, (float(value),) * out_dim)
    return HardConstraintField(net=net, mask=mask, lift=lift)

Related APIsomnibias.jax.architectures.AffineFactor↗omnibias.jax.architectures.AffineLift↗omnibias.jax.architectures.BoundaryMask↗omnibias.jax.architectures.HardConstraintField↗

Other public import paths

• omnibias.jax.architectures.initial_value
• omnibias.jax.architectures.hardbc.initial_value

← Package overviewRunnable learning examples →