Backends

A backend is a thin specialization of omnibias-core for one tensor library. Each provides the same activation math; they differ only in tensor type and framework-native layers.

PyTorch (omnibias-torch)

The most feature-complete backend.

• OMBU — the trainable K-bias operator unit.
• OperatorBlock — the typed scalar operator.
• cmbLinear, cmbConv* — nn.Linear / nn.Conv* with an inline operator.
• GrowableOMBU — increase K during training under a scheduler.
• Reference architectures: PINNHeat, CmbNet, CvxLasso, CvxLogistic.
• Fast-path kernels and training schedulers.

from omnibias.torch import OMBU, OperatorBlock, cmbLinear

JAX (omnibias-jax)

The numerics-first backend.

• The activation dictionary (real and complex).
• Closed-form Laplacian / Hessian for one-layer fields.
• neural_field_value_grad_laplacian, neural_field_polylaplacian.
• Born–Oppenheimer derivative kernels for VMC.
• The jet / jet_mv kernels for directional and multivariate jets.

from omnibias.jax import get_activation, neural_field_value_grad_hessian

Keras 3 (omnibias-keras)

A unified backend: the same code runs on TensorFlow, JAX, or PyTorch via Keras's backend selector. Choose with KERAS_BACKEND before importing keras.

• OMBU, OperatorBlock, cmbDense.
• Activation-level math is bit-identical to the other backends.

import os; os.environ["KERAS_BACKEND"] = "jax"
import keras
from omnibias.keras import OMBU, cmbDense

FermiNet bridge (omnibias-ferminet)

Depends on omnibias-jax. Provides a folx-compatible Laplacian and the production local-kinetic-energy factories for neural VMC. See the VMC tutorial.

How bit-parity is maintained

Because all three import the same coefficients, a dedicated parity suite can assert float64-ULP equality on every (activation, order) pair. See Cross-backend parity.

What lives where

Need	Backend module
PyTorch layers	omnibias.torch.unit, .blocks, .growable
JAX closed-form fields	omnibias.jax.laplacian
Keras unified layers	omnibias.keras
FermiNet local kinetic energy	omnibias.ferminet.integration

See the API reference for the full per-package surface.