omnibias.jax
The JAX backend. Bit-identical to omnibias.torch because both import the same
core coefficients. Install with pip install omnibias-jax.
import omnibias.jax as oj
Closed-form neural-field kernels
The closed-form operators for the one-layer field
f(x) = b + Σ_h c_h · σ(W_h·x + b_h) on ℝᴰ.
| Symbol | Returns |
|---|---|
neural_field_value | f(x) |
neural_field_laplacian | Δf(x) |
neural_field_hessian | ∇²f(x) |
neural_field_value_and_laplacian | (f, Δf) |
neural_field_value_grad_laplacian | (f, ∇f, Δf) |
neural_field_value_grad_hessian | (f, ∇f, ∇²f) |
These are the kernels behind the first Laplacian tutorial.
Born–Oppenheimer derivatives
Used for nuclear-Hessian / molecular work.
coulomb_potential · make_local_energy · make_bo_force ·
make_bo_hessian · vibrational_frequencies
Activation registry
get_activation · list_activations · register_activation ·
is_registered · JaxActivationSpec (alias of ActivationSpec pinned to
jax.Array).
Jets
- directional:
mlp_jet,layer_jet,affine_jet,compose_jet,tower_to_jet,jet_to_tower,lhopital_ratio,limit_of_ratio,removable_value. - multivariate:
mlp_jet_mv,layer_jet_mv,affine_jet_mv,compose_jet_mv,identity_jet,jet_multiply,jet_partials,jet_gradient,jet_hessian.
Statistical surfaces
Mirrors the torch backend: divergences (kl_divergence, js_divergence,
renyi_divergence, wasserstein1, wasserstein2_gaussian, …), GLM helpers
(glm_mean, glm_variance, exponential_family_cumulants,
fit_natural_parameter, moment_match), moments
(gaussian_moment_propagation, delta_method_*), and probability (cdf,
soft_histogram, ks_statistic, binned_calibration_error,
empirical_band_mass, model_band_mass).
The iterated neural_field_polylaplacian (Δᵏ) used in the
high-order PDEs guide lives alongside these
kernels in the JAX field surface.