omnibias.jax.confluence

0.5.0rc1Release candidatePython >=3.10Apache-2.0

JAX activation derivatives, functional network jets, direct contractions, and parameter-differentiable numerical models.

import omnibias.jax.confluence

Module contract & mathematical conventions

Stable centered collision atoms, live through the attained rho=0 boundary.

Finite formulas are exact in real arithmetic. For abs(h*eta)<=series_radius an explicit even Taylor polynomial is evaluated (series_terms coefficients). Its remainder is bounded by the shared core confluence derivative bounds. The inactive finite branch uses positive denominators so backward/JIT are safe.

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

CenteredPair

omnibias.jax.confluence.CenteredPairSource ↗

class CenteredPair(NamedTuple)

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Fields & methods

averagefield

average: Array

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

View source ↗

divided_differencefield

divided_difference: Array

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

View source ↗

Read the implementation

class CenteredPair(NamedTuple):
    average: Array
    divided_difference: Array

FUNCTION

centered_pair

omnibias.jax.confluence.centered_pairSource ↗

centered_pair(z: Array, rho: Array, eta: Array, *, activation: str='sigmoid', series_terms: int=6, series_radius: float=0.01, tower: Array | Sequence[Array] | None=None) -> CenteredPair

A=(sigma(z+h*eta)+sigma(z-h*eta))/2, B=odd difference/(2h).

rho=h**2 is nonnegative (projected in this function). At rho=0, A=sigma(z) and B=eta*sigma'(z). Includes live derivatives in rho, z, and eta. Sigmoid and tanh have stable finite-spread formulas; other activations are rejected. The integer series_terms and radius are static compilation parameters.

Parameters

Name	Type & default	Contract
`z`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`rho`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`eta`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`activation`keyword-only	`str`Default: 'sigmoid'	See the operation contract and implementation below.
`series_terms`keyword-only	`int`Default: 6	See the operation contract and implementation below.
`series_radius`keyword-only	`float`Default: 0.01	See the operation contract and implementation below.
`tower`keyword-only	`Array \| Sequence[Array] \| None`Default: None	See the operation contract and implementation below.

ReturnsCenteredPair

Explicit validation & failure cases

When activation not in ('sigmoid', 'tanh'), raises ValueError('centered_pair supports sigmoid and tanh').
When series_terms < 1 or not 0 < series_radius <= 0.1, raises ValueError('positive terms and radius in (0,0.1] required').
When spec.fastpath is None, raises ValueError('activation requires a derivative provider').
When len(tower) <= maximum, raises ValueError('supplied tower is too short for the centered series').

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

Regression example · values and live derivatives

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_values_and_live_derivatives(activation, rho):
    z = torch.tensor([-2.1, -0.3, 0.4, 2.0], dtype=torch.float64, requires_grad=True)
    p = torch.tensor(rho, dtype=torch.float64, requires_grad=True)
    eta = torch.tensor([-1.4, -0.2, 0.7, 1.2], dtype=torch.float64, requires_grad=True)
    a, b = tpair(z, p, eta, activation=activation)
    ja, jb = jax.jit(lambda x, r, e: jpair(x, r, e, activation=activation))(
        jnp.array(z.detach()), jnp.array(rho), jnp.array(eta.detach())
    )
    np.testing.assert_allclose(a.detach(), ja, rtol=1e-10, atol=1e-12)
    np.testing.assert_allclose(b.detach(), jb, rtol=1e-10, atol=1e-12)
    act = torch.sigmoid if activation == "sigmoid" else torch.tanh
    if rho >= 0.03:
        h = math.sqrt(rho)
        expected_a = (act(z + h * eta) + act(z - h * eta)) / 2
        expected_b = (act(z + h * eta) - act(z - h * eta)) / (2 * h)
    else:
        # Independent differentiable activation Taylor expansion, with AD
        # derivatives rather than the shared integer coefficient evaluator.
        fn = act
        derivs = [fn(z)]
        for _ in range(11):
            derivs.append(torch.autograd.grad(derivs[-1].sum(), z, create_graph=True)[0])
        expected_a = sum(
            derivs[2 * k] * rho**k * eta ** (2 * k) / math.factorial(2 * k) for k in range(6)
        )
        expected_b = sum(
            derivs[2 * k + 1] * rho**k * eta ** (2 * k + 1) / math.factorial(2 * k + 1)
            for k in range(6)
        )
    torch.testing.assert_close(a, expected_a, rtol=1e-10, atol=1e-12)
    torch.testing.assert_close(b, expected_b, rtol=1e-10, atol=1e-12)
    actual = a + b
    expected = expected_a + expected_b
    for _ in range(4):
        actual = torch.autograd.grad(actual.sum(), z, create_graph=True)[0]
        expected = torch.autograd.grad(expected.sum(), z, create_graph=True)[0]
        torch.testing.assert_close(actual, expected, rtol=1e-10, atol=1e-12)
    gp = torch.autograd.grad((a + b).sum(), p)[0]
    jp = jax.grad(
        lambda r: sum(
            jpair(jnp.array(z.detach()), r, jnp.array(eta.detach()), activation=activation)
        ).sum()
    )(jnp.array(rho))
    np.testing.assert_allclose(gp.detach(), jp, rtol=1e-10, atol=1e-12)

python -m pytest tests/test_confluence_parity.py -k test_values_and_live_derivatives -q

Read the implementation

def centered_pair(
    z: Array, rho: Array, eta: Array, *, activation: str = "sigmoid",
    series_terms: int = 6, series_radius: float = 0.01,
    tower: Array | Sequence[Array] | None = None,
) -> CenteredPair:
    """A=(sigma(z+h*eta)+sigma(z-h*eta))/2, B=odd difference/(2h).

    rho=h**2 is nonnegative (projected in this function). At rho=0, A=sigma(z)
    and B=eta*sigma'(z). Includes live derivatives in rho, z, and eta. Sigmoid
    and tanh have stable finite-spread formulas; other activations are rejected.
    The integer series_terms and radius are static compilation parameters.
    """
    if activation not in ("sigmoid", "tanh"):
        raise ValueError("centered_pair supports sigmoid and tanh")
    if series_terms < 1 or not 0 < series_radius <= 0.1:
        raise ValueError("positive terms and radius in (0,0.1] required")
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError("activation requires a derivative provider")
    p = retract_spread(rho)
    small = p * eta**2 <= series_radius**2
    safe_rho = jnp.where(small, jnp.ones_like(p), p)
    h = jnp.sqrt(safe_rho)
    u = jnp.abs(h * eta)
    u = jnp.where(small, jnp.ones_like(u), u)
    # Three sigmoid evaluations in the finite path; no subtraction of nearby
    # activation values. tanh(z)=2*sigmoid(2*z)-1 supplies the same formula.
    scale = 1.0 if activation == "sigmoid" else 2.0
    zz = scale * z
    uu = scale * u
    splus = jax.nn.sigmoid(zz + uu)
    sminus = jax.nn.sigmoid(zz - uu)
    complement = jax.nn.sigmoid(-zz + uu)
    a_finite = (splus + sminus) * 0.5
    b_finite = eta * scale * splus * complement * (-jnp.expm1(-2 * uu) / (2 * uu))
    if activation == "tanh":
        a_finite = 2 * a_finite - 1
        b_finite = 2 * b_finite
    # Evaluate the near branch at safe spread even when it is inactive: huge
    # finite spreads must not overflow powers in a discarded Taylor polynomial.
    r = jnp.where(small, p, jnp.zeros_like(p))
    e = jnp.where(small, eta, jnp.zeros_like(eta))
    maximum = 2 * series_terms - 1
    if tower is None:
        tower = (spec.tower(z, maximum) if spec.tower is not None
                 else tuple(spec.fastpath(z, k) for k in range(maximum + 1)))
    if len(tower) <= maximum:
        raise ValueError("supplied tower is too short for the centered series")
    a_series = jnp.zeros_like(z)
    b_series = jnp.zeros_like(z)
    for k in range(series_terms):
        a_series = a_series + tower[2*k] * (r**k * e**(2*k)) / factorial(2*k)
        b_series = b_series + tower[2*k+1] * (r**k * e**(2*k+1)) / factorial(2*k+1)
    return CenteredPair(jnp.where(small, a_series, a_finite),
                        jnp.where(small, b_series, b_finite))

FUNCTION

collision_atom

omnibias.jax.confluence.collision_atomSource ↗

collision_atom(z: Array, rho: Array, eta: Array, m0: Array, m1: Array, *, activation: str='sigmoid', series_terms: int=6, tower: Array | Sequence[Array] | None=None) -> Array

m0*A + m1*B; moments remain bounded as a pair becomes a derivative atom.

Parameters

Name	Type & default	Contract
`z`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`rho`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`eta`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`m0`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`m1`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`activation`keyword-only	`str`Default: 'sigmoid'	See the operation contract and implementation below.
`series_terms`keyword-only	`int`Default: 6	See the operation contract and implementation below.
`tower`keyword-only	`Array \| Sequence[Array] \| None`Default: None	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def collision_atom(z: Array, rho: Array, eta: Array, m0: Array, m1: Array,
                   *, activation: str = "sigmoid", series_terms: int = 6,
                   tower: Array | Sequence[Array] | None = None) -> Array:
    """m0*A + m1*B; moments remain bounded as a pair becomes a derivative atom."""
    pair = centered_pair(z, rho, eta, activation=activation, series_terms=series_terms, tower=tower)
    return m0 * pair.average + m1 * pair.divided_difference

FUNCTION

moment_atom

omnibias.jax.confluence.moment_atomSource ↗

moment_atom(z: Array, moments: Array, *, activation: str='sigmoid', tower: Array | Sequence[Array] | None=None) -> Array

A truncated common-weight bias cluster, sum(m[k]*sigma**(k)(z)).

The host initialize_moments/cluster_remainder API supplies its error bound. This is exact only when the selected atom architecture is itself the model.

Parameters

Name	Type & default	Contract
`z`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`moments`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`activation`keyword-only	`str`Default: 'sigmoid'	See the operation contract and implementation below.
`tower`keyword-only	`Array \| Sequence[Array] \| None`Default: None	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When moments.ndim != 1 or moments.shape[0] < 1, raises ValueError('moments must be a nonempty vector').
When spec.fastpath is None, raises ValueError('activation requires a derivative provider').
When len(tower) <= n, raises ValueError('supplied tower is too short for the moment expansion').

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

Read the implementation

def moment_atom(z: Array, moments: Array, *, activation: str = "sigmoid",
                tower: Array | Sequence[Array] | None = None) -> Array:
    """A truncated common-weight bias cluster, sum(m[k]*sigma**(k)(z)).

    The host initialize_moments/cluster_remainder API supplies its error bound.
    This is exact only when the selected atom architecture is itself the model.
    """
    if moments.ndim != 1 or moments.shape[0] < 1:
        raise ValueError("moments must be a nonempty vector")
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError("activation requires a derivative provider")
    n = moments.shape[0] - 1
    if tower is None:
        tower = (spec.tower(z, n) if spec.tower is not None
                 else tuple(spec.fastpath(z, k) for k in range(n+1)))
    if len(tower) <= n:
        raise ValueError("supplied tower is too short for the moment expansion")
    out = jnp.zeros_like(z)
    for k in range(n+1):
        out = out + moments[k] * tower[k]
    return out

FUNCTION

polynomial_pair

omnibias.jax.confluence.polynomial_pairSource ↗

polynomial_pair(z: Array, rho: Array, eta: Array, coefficients: Array) -> CenteredPair

Exact finite polynomial control in the same centered coordinates.

Coefficients are live and in ascending powers. No finite-spread truncation occurs: the even/odd binomial expansion terminates at the supplied degree.

Parameters

Name	Type & default	Contract
`z`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`rho`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`eta`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`coefficients`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsCenteredPair

Explicit validation & failure cases

When coefficients.ndim != 1 or coefficients.shape[0] < 1, raises ValueError('nonempty polynomial coefficient vector required').

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

Regression example · polynomial pair rejects invalid coefficient shape

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_polynomial_pair_rejects_invalid_coefficient_shape():
    for coefficients in (np.empty(0), np.ones((2, 2))):
        with pytest.raises(ValueError, match="coefficient"):
            torch_pair(
                torch.tensor(0.0), torch.tensor(0.0), torch.tensor(1.0), torch.tensor(coefficients)
            )
        with pytest.raises(ValueError, match="coefficient"):
            jax_pair(
                jnp.asarray(0.0), jnp.asarray(0.0), jnp.asarray(1.0), jnp.asarray(coefficients)
            )

python -m pytest tests/test_polynomial_pair.py -k test_polynomial_pair_rejects_invalid_coefficient_shape -q

Regression example · polynomial pair values live coefficients and spread derivatives

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_polynomial_pair_values_live_coefficients_and_spread_derivatives(rho_value, coefficients):
    z = np.array([-0.75, -0.25, 0.0, 0.5])
    eta = np.array([0.5, -0.75, 0.0, 1.25])
    coefficients = np.asarray(coefficients)
    expected_a, expected_b, da, db = oracle(coefficients, z, rho_value, eta)
    coeff = torch.tensor(coefficients, requires_grad=True)
    rho = torch.tensor(rho_value, dtype=torch.float64, requires_grad=True)
    ta, tb = torch_pair(torch.tensor(z), rho, torch.tensor(eta), coeff)
    loss = (ta + 0.25 * tb).sum()
    dc, drho = torch.autograd.grad(loss, (coeff, rho))

    def jloss(c, r):
        a, b = jax_pair(jnp.asarray(z), r, jnp.asarray(eta), c)
        return (a + 0.25 * b).sum()

    ja, jb = jax_pair(
        jnp.asarray(z), jnp.asarray(rho_value), jnp.asarray(eta), jnp.asarray(coefficients)
    )
    jdc, jdrho = jax.jit(jax.grad(jloss, argnums=(0, 1)))(
        jnp.asarray(coefficients), jnp.asarray(rho_value)
    )
    ca, cb = jax.jit(jax_pair)(
        jnp.asarray(z), jnp.asarray(rho_value), jnp.asarray(eta), jnp.asarray(coefficients)
    )
    for actual, expected in (
        (ta.detach(), expected_a),
        (tb.detach(), expected_b),
        (ja, expected_a),
        (jb, expected_b),
        (ca, expected_a),
        (cb, expected_b),
    ):
        np.testing.assert_allclose(actual, expected, atol=2e-15, rtol=2e-15)
    np.testing.assert_array_equal(ta.detach(), ja)
    np.testing.assert_array_equal(tb.detach(), jb)
    expected_coeff = []
    for basis in np.eye(len(coefficients)):
        a, b, _, _ = oracle(basis, z, rho_value, eta)
        expected_coeff.append((a + 0.25 * b).sum())
    for actual in (dc, jdc):
        np.testing.assert_allclose(actual, expected_coeff, atol=2e-15, rtol=2e-15)
    for actual in (drho, jdrho):
        np.testing.assert_allclose(actual, (da + 0.25 * db).sum(), atol=2e-15, rtol=2e-15)

python -m pytest tests/test_polynomial_pair.py -k test_polynomial_pair_values_live_coefficients_and_spread_derivatives -q

Read the implementation

def polynomial_pair(z: Array, rho: Array, eta: Array, coefficients: Array) -> CenteredPair:
    """Exact finite polynomial control in the same centered coordinates.

    Coefficients are live and in ascending powers. No finite-spread truncation
    occurs: the even/odd binomial expansion terminates at the supplied degree.
    """
    from math import comb
    if coefficients.ndim != 1 or coefficients.shape[0] < 1:
        raise ValueError('nonempty polynomial coefficient vector required')
    p = retract_spread(rho)
    a, b = jnp.zeros_like(z), jnp.zeros_like(z)
    for degree in range(coefficients.shape[0]):
        for power in range(degree+1):
            term = coefficients[degree]*comb(degree,power)*z**(degree-power)
            if power % 2:
                b = b+term*p**((power-1)//2)*eta**power
            else:
                a = a+term*p**(power//2)*eta**power
    return CenteredPair(a,b)

FUNCTION

retract_spread

omnibias.jax.confluence.retract_spreadSource ↗

retract_spread(rho: Array) -> Array

Project onto rho>=0; call between steps to attain zero exactly.

Parameters

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

ReturnsArray

Regression example · jit parameter gradient retraction and moment mode

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_jit_parameter_gradient_retraction_and_moment_mode() -> None:
    _, bank = pair_banks()
    bank, _, _ = jb.commit_transition(bank, bank.propose_pair(0, 1), moments(bank), between_steps=True,
                                      acceptance=lambda _b, _p: True)
    x = jnp.asarray([-0.3, 0.2])
    compiled = jax.jit(lambda state, z: state.forward(z))
    np.testing.assert_array_equal(np.asarray(compiled(bank, x)), np.asarray(bank.forward(x)))
    def objective(parameters):
        return jnp.sum(bank.with_parameters(parameters).forward(x) ** 2)
    gradients = jax.jit(jax.grad(objective))(bank.parameters)
    assert all(bool(jnp.all(jnp.isfinite(value))) for value in gradients.values())
    assert float(jnp.abs(gradients["odd_moments"][0])) > 0
    boundary = replace(bank, rho=bank.rho.at[0].set(-0.5)).retract()
    assert float(boundary.rho[0]) == 0
    assert float(jax.grad(lambda r: jb.retract_spread(r))(jnp.asarray(0.0))) == 1
    moment_bank = replace(boundary, modes=boundary.modes.at[0].set(2),
                          moment_coefficients=boundary.moment_coefficients.at[0, :2].set(jnp.array([0.7, 0.2])))
    assert bool(jnp.all(jnp.isfinite(compiled(moment_bank, x))))

python -m pytest tests/test_confluent_bank_parity.py -k test_jit_parameter_gradient_retraction_and_moment_mode -q

Read the implementation

def retract_spread(rho: Array) -> Array:
    """Project onto rho>=0; call between steps to attain zero exactly."""
    return jnp.where(rho >= 0, rho, jnp.zeros_like(rho))

← Package overview Runnable learning examples →

omnibias.jax.confluence

CenteredPair​

Fields & methods

centered_pair​

Parameters

collision_atom​

Parameters

moment_atom​

Parameters

polynomial_pair​

Parameters

retract_spread​

Parameters

CenteredPair

centered_pair

collision_atom

moment_atom

polynomial_pair

retract_spread