Skip to main content
API REFERENCE/omnibias-jax

omnibias.jax.precision

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.precision
Module contract & mathematical conventions

The precision precondition behind every "algebraically corresponding twin" claim.

omnibias computes the same closed-form derivative tower on PyTorch and JAX from one shared coefficient module, so the two backends evaluate corresponding formulas; use double precision when comparing their numerical results. PyTorch uses float64 for Python floats; JAX silently truncates arrays to float32 unless 64-bit mode is enabled. Under the JAX default the twins are each internally consistent but agree only to float32 tolerance.

That gap is not always cosmetic. Wherever a value feeds a threshold, a rounding step or an argmax -- the decoders in omnibias.discrete / omnibias.qubo / omnibias.struct, the hardening step in omnibias.partition -- a 1e-7 disagreement can flip a bit and move the answer by 1, so the divergence shows up as a different decision rather than a different last digit.

This module deliberately does **not** flip the flag for you. jax_enable_x64 is process-global and cannot be changed once arrays exist, so a library that sets it on import would silently re-specify the numerics of every other piece of JAX code in the process. Enable it yourself, before the first JAX array::

import jax jax.config.update("jax_enable_x64", True)

or set JAX_ENABLE_X64=1 in the environment.

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

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →
CONSTANT

X64_HINT​

omnibias.jax.X64_HINTSource ↗
X64_HINT = 'enable 64-bit JAX before the first array is created -- jax.config.update("jax_enable_x64", True) or JAX_ENABLE_X64=1'

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Regression example · the precision helper reports the live flag

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_the_precision_helper_reports_the_live_flag() -> None:
    import pytest

    pytest.importorskip("jax")
    import jax
    from omnibias.jax.precision import X64_HINT, require_x64, x64_enabled

    assert x64_enabled() is bool(jax.config.jax_enable_x64)
    assert "jax_enable_x64" in X64_HINT
    if x64_enabled():
        require_x64()  # must not raise
    else:
        with pytest.raises(RuntimeError, match="64-bit JAX"):
            require_x64()
python -m pytest tests/test_bit_identity_precondition.py -k test_the_precision_helper_reports_the_live_flag -q
Read the implementation
X64_HINT = (
    'enable 64-bit JAX before the first array is created -- '
    'jax.config.update("jax_enable_x64", True) or JAX_ENABLE_X64=1'
)
Other public import paths
  • omnibias.jax.X64_HINT
  • omnibias.jax.precision.X64_HINT
FUNCTION

require_x64​

omnibias.jax.require_x64Source ↗
require_x64(context: str='bit-identical torch/jax parity') -> None

Raise unless 64-bit JAX is on.

Call this at the top of a script or test whose correctness depends on the parity claim, so the precondition fails loudly instead of degrading into a silent float32 comparison.

Parameters

NameType & defaultContract
contextpositional-or-keywordstrDefault: 'bit-identical torch/jax parity'See the operation contract and implementation below.
ReturnsNone
Explicit validation & failure cases
  • When not x64_enabled(), raises RuntimeError(f'{context} requires 64-bit JAX: {X64_HINT}').

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

Regression example · the precision helper reports the live flag

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_the_precision_helper_reports_the_live_flag() -> None:
    import pytest

    pytest.importorskip("jax")
    import jax
    from omnibias.jax.precision import X64_HINT, require_x64, x64_enabled

    assert x64_enabled() is bool(jax.config.jax_enable_x64)
    assert "jax_enable_x64" in X64_HINT
    if x64_enabled():
        require_x64()  # must not raise
    else:
        with pytest.raises(RuntimeError, match="64-bit JAX"):
            require_x64()
python -m pytest tests/test_bit_identity_precondition.py -k test_the_precision_helper_reports_the_live_flag -q
Read the implementation
def require_x64(context: str = "bit-identical torch/jax parity") -> None:
    """Raise unless 64-bit JAX is on.

    Call this at the top of a script or test whose correctness depends on the
    parity claim, so the precondition fails loudly instead of degrading into a
    silent ``float32`` comparison.
    """
    if not x64_enabled():
        raise RuntimeError(f"{context} requires 64-bit JAX: {X64_HINT}")
Related APIsomnibias.jax.x64_enabled↗
Other public import paths
  • omnibias.jax.require_x64
  • omnibias.jax.precision.require_x64
FUNCTION

x64_enabled​

omnibias.jax.x64_enabledSource ↗
x64_enabled() -> bool

Whether JAX is in 64-bit mode, i.e. whether bit-parity is achievable.

Returnsbool
Regression example · the precision helper reports the live flag

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_the_precision_helper_reports_the_live_flag() -> None:
    import pytest

    pytest.importorskip("jax")
    import jax
    from omnibias.jax.precision import X64_HINT, require_x64, x64_enabled

    assert x64_enabled() is bool(jax.config.jax_enable_x64)
    assert "jax_enable_x64" in X64_HINT
    if x64_enabled():
        require_x64()  # must not raise
    else:
        with pytest.raises(RuntimeError, match="64-bit JAX"):
            require_x64()
python -m pytest tests/test_bit_identity_precondition.py -k test_the_precision_helper_reports_the_live_flag -q
Read the implementation
def x64_enabled() -> bool:
    """Whether JAX is in 64-bit mode, i.e. whether bit-parity is achievable."""
    import jax

    return bool(getattr(jax.config, "jax_enable_x64", False))
Other public import paths
  • omnibias.jax.x64_enabled
  • omnibias.jax.precision.x64_enabled
← Package overviewRunnable learning examples →