Metadata-Version: 2.4
Name: omnibias-jax
Version: 0.4.0
Summary: JAX backend for omnibias: closed-form n-th derivative activation kernels (sigmoid via Eulerian polynomials, tanh via Legendre, Gaussian via Hermite), neural-field Laplacian / Hessian primitives, and Born-Oppenheimer derivative tools for variational quantum Monte Carlo.
Author-email: Vardan Grigoryants <vardan@derivon.ai>
Maintainer-email: Derivon <info@derivon.ai>
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/derivon-ai/omnibias
Project-URL: Documentation, https://omnibias.ai/
Project-URL: Source, https://github.com/derivon-ai/omnibias
Project-URL: Issues, https://github.com/derivon-ai/omnibias/issues
Project-URL: Changelog, https://github.com/derivon-ai/omnibias/blob/main/CHANGELOG.md
Keywords: deep-learning,jax,activation-functions,neural-quantum-states,variational-monte-carlo,physics-informed-neural-networks,closed-form-derivatives,born-oppenheimer
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Physics
Classifier: Typing :: Typed
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: omnibias-core>=0.2.0
Requires-Dist: jax>=0.4.30
Requires-Dist: jaxlib>=0.4.30
Requires-Dist: numpy>=1.23
Provides-Extra: test
Requires-Dist: pytest>=7.0; extra == "test"
Provides-Extra: examples
Requires-Dist: matplotlib>=3.6; extra == "examples"
Requires-Dist: scipy>=1.10; extra == "examples"
Dynamic: license-file

# omnibias-jax

JAX backend for the omnibias closed-form n-th derivative framework.

## Why this is fast

All numbers float64, identical answers to autodiff up to `≤ 10⁻¹⁵`. Full
derivation in [`docs/complexity.md`](../../docs/complexity.md).

- Closed-form Laplacian overhead is **`O(1)` in input dimension `D`** —
  independent of `D` because the reduction collapses the inner sum once.
- At `D = 240`, **68× faster than `jax.hessian` + trace** and **63× less
  memory**.
- Iterated Laplacian `Δᵏ` is **480× faster than folx-nested at `k = 3`**;
  folx-nested OOMs at `k = 4` while omnibias finishes in ~0.1 ms.
- Bit-identical to `omnibias-torch` and `omnibias-keras` for every
  `(activation, order)` pair.

## Install

```bash
pip install omnibias-jax
```

`omnibias-jax` depends on `omnibias-core` (pure-Python math) and `jax>=0.4.30`.

## What is in here

- The same closed-form derivative kernels as `omnibias-torch`, written in
  JAX (`jax.numpy`) so they JIT-compile cleanly inside FermiNet, vmc_jax,
  DeepQMC, and similar stacks. Polynomial coefficients are imported from
  `omnibias-core`, so a JAX `sigma^(n)(z)` is *bit-identical* to the
  torch `sigma^(n)(z)` for every `n` (this is the contract validated by
  `tests/test_cross_backend_parity.py`).
- A backend-specific activation dictionary registered via the same
  ``ActivationSpec`` protocol as torch (`get_activation`, `list_activations`,
  `register_activation`).
- ``neural_field_laplacian`` / ``neural_field_hessian`` / family: closed-form
  Laplacian and full Hessian for a one-layer scalar field
  ``f(x) = b + sum_h c_h sigma(W_h . x + b_h)`` on ``R^D``. These are the
  primitives the FermiNet bridge in `omnibias-ferminet` calls when composing
  through coordinate transformations.
- Born-Oppenheimer derivative kernels (`coulomb_potential`,
  `make_local_energy`, `make_bo_force`, `make_bo_hessian`,
  `vibrational_frequencies`) used to build analytic nuclear Hessians of
  neural-VMC energies.

## Public API

```python
from omnibias.jax import (
    JaxActivationSpec, get_activation, list_activations,
    register_activation, is_registered,
    neural_field_laplacian, neural_field_value_grad_hessian,
    coulomb_potential, make_local_energy,
    make_bo_force, make_bo_hessian, vibrational_frequencies,
    BankSpec, init_bias_scan, bias_scan, init_multipack, multipack_apply,
)
```

The FermiNet bridge (`folx`-compatible API, Tier-2 restricted FermiNet,
multiblock primitives) lives in the separate `omnibias-ferminet`
package; importing `omnibias.jax` does **not** trigger a FermiNet import,
so the JAX core remains useful when FermiNet is absent.

Wave-1 twins: `init_multipack` / `multipack_apply` (01-01, **shipped**)
and `init_bias_scan` / `bias_scan` / `BankSpec` (01-02, **shipped**). Same honesty as the
torch modules: interior shift along `w`, `gamma` is not `delta -> 0`.

Shipped Wave-3 twins: Scan-Net (`init_scan_net` / `scan_net_apply`;
on-lattice equivariance; G4 leftover-recorded), Jet-KAN (`init_jet_kan`
/ `jet_kan_apply`; model-jet exactness, KA theorem does not justify; G2
leftover-recorded), equivariant scan (gaussian-family steering;
discrete `C_L`; G5 leftover-recorded), `hierarchical_scan` (1-D
offsets), and Hermite ladder (`hermite_basis` / `ladder_apply`;
Rodrigues reweight required; G4 leftover-recorded).

## License

Apache-2.0. See [`LICENSE`](LICENSE) and [`../../LICENSING.md`](../../LICENSING.md).
You never need a commercial licence for this package.
