omnibias.jax
JAX activation derivatives, functional network jets, direct contractions, and parameter-differentiable numerical models.
import omnibias.jaxUse this backend with JAX arrays and functional parameter trees. The forward derivative calculation is distinct from the JAX transformation you use to differentiate a loss with respect to model parameters. A high spatial order does not by itself require nesting that many reverse-mode spatial transformations.
Enable JAX x64 before constructing arrays if your validation requires float64. Treat compilation and steady-state evaluation as separate benchmark phases. Derivative order, architecture, and other structural arguments may need to be static under jit; follow the selected function's contract rather than assuming every Python option is traceable.
Directional jets, mixed-partial jets, and direct Laplacian contractions have different output layouts. Keep the factorial convention explicit when moving between Taylor coefficients and derivatives. The JAX high-order-residual and Laplacian examples demonstrate value_and_grad through the chosen spatial path. Shared formulas across backends do not imply universal bitwise identity.
Module contract & mathematical conventions
JAX activation derivatives and Taylor jets for physics-informed models.
Shared polynomial coefficients define the derivative tower. Operator blocks, mixed jets, and PINN architectures reuse it without nested coordinate autodiff. Architecture and optimization helpers are available in their own modules.
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.
activationsModule ↗activations_complexModule ↗architecturesModule ↗bo_derivativesModule ↗confluenceModule ↗confluent_bankModule ↗implicitModule ↗jetModule ↗jet_mvModule ↗laplacianModule ↗line_searchModule ↗multipackModule ↗optimModule ↗optim_block_searchModule ↗precisionModule ↗realizationModule ↗refineModule ↗scanModule ↗scan_equivariantModule ↗weight_loss_jetModule ↗Imports & aliases
These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.
| Import from this module | Canonical reference |
|---|---|
AdaptivePackBank | omnibias.jax.AdaptivePackBank |
BankSpec | omnibias.core.BankSpec |
BirkhoffOMBU | omnibias.jax.BirkhoffOMBU |
BlockSpec | omnibias.core.BlockSpec |
DEQConfig | omnibias.core.DEQConfig |
DEQNotContractive | omnibias.core.DEQNotContractive |
DEQResult | omnibias.jax.DEQResult |
DEQSolverUnknown | omnibias.core.DEQSolverUnknown |
GradientSecant | omnibias.core.GradientSecant |
JaxActivationSpec | omnibias.jax.JaxActivationSpec |
JetLineSearchConfig | omnibias.core.JetLineSearchConfig |
LineSearchResult | omnibias.core.LineSearchResult |
WeightLossJetSpec | omnibias.core.WeightLossJetSpec |
X64_HINT | omnibias.jax.X64_HINT |
affine_jet | omnibias.jax.affine_jet |
affine_jet_mv | omnibias.jax.affine_jet_mv |
antiderivative_jet | omnibias.jax.antiderivative_jet |
arrangement_w_block | omnibias.core.arrangement_w_block |
bank_forward | omnibias.jax.bank_forward |
bias_scan | omnibias.jax.bias_scan |
block_direction | omnibias.jax.block_direction |
block_exact_search | omnibias.jax.block_exact_search |
block_exact_sweep | omnibias.jax.block_exact_sweep |
compose_jet | omnibias.jax.compose_jet |
compose_jet_mv | omnibias.jax.compose_jet_mv |
compose_jet_riccati | omnibias.jax.compose_jet_riccati |
coulomb_potential | omnibias.jax.coulomb_potential |
deq_du_dW | omnibias.jax.deq_du_dW |
deq_solve | omnibias.jax.deq_solve |
deq_vjp | omnibias.jax.deq_vjp |
derivative_jet | omnibias.jax.derivative_jet |
equivariant_scan_apply | omnibias.jax.equivariant_scan_apply |
get_activation | omnibias.jax.get_activation |
identity_jet | omnibias.jax.identity_jet |
init_bias_scan | omnibias.jax.init_bias_scan |
init_multipack | omnibias.jax.init_multipack |
init_pack_bank | omnibias.jax.init_pack_bank |
is_registered | omnibias.jax.is_registered |
jet_attention | omnibias.jax.jet_attention |
jet_exp | omnibias.jax.jet_exp |
jet_gradient | omnibias.jax.jet_gradient |
jet_hessian | omnibias.jax.jet_hessian |
jet_line_search | omnibias.jax.jet_line_search |
jet_line_search_on_ray | omnibias.jax.jet_line_search_on_ray |
jet_multiply | omnibias.jax.jet_multiply |
jet_partials | omnibias.jax.jet_partials |
jet_reciprocal | omnibias.jax.jet_reciprocal |
jet_softmax | omnibias.jax.jet_softmax |
jet_to_tower | omnibias.jax.jet_to_tower |
last_linear_block | omnibias.core.last_linear_block |
layer_jet | omnibias.jax.layer_jet |
layer_jet_mv | omnibias.jax.layer_jet_mv |
lhopital_ratio | omnibias.jax.lhopital_ratio |
limit_of_ratio | omnibias.jax.limit_of_ratio |
list_activations | omnibias.jax.list_activations |
make_bo_force | omnibias.jax.make_bo_force |
make_bo_hessian | omnibias.jax.make_bo_hessian |
make_local_energy | omnibias.jax.make_local_energy |
mlp_jet | omnibias.jax.mlp_jet |
mlp_jet_mv | omnibias.jax.mlp_jet_mv |
multipack_apply | omnibias.jax.BirkhoffOMBU |
multipack_response | omnibias.jax.multipack_response |
neural_field_hessian | omnibias.jax.neural_field_hessian |
neural_field_laplacian | omnibias.jax.neural_field_laplacian |
neural_field_value | omnibias.jax.neural_field_value |
neural_field_value_and_laplacian | omnibias.jax.neural_field_value_and_laplacian |
neural_field_value_grad_hessian | omnibias.jax.neural_field_value_grad_hessian |
neural_field_value_grad_laplacian | omnibias.jax.neural_field_value_grad_laplacian |
ombu_bias_block | omnibias.core.ombu_bias_block |
one_layer_loss | omnibias.jax.one_layer_loss |
one_layer_loss_grad | omnibias.jax.one_layer_loss_grad |
one_layer_loss_hessian | omnibias.jax.one_layer_loss_hessian |
one_layer_loss_jet | omnibias.jax.one_layer_loss_jet |
one_layer_newton_direction | omnibias.jax.one_layer_newton_direction |
register_activation | omnibias.jax.register_activation |
removable_value | omnibias.jax.removable_value |
require_x64 | omnibias.jax.require_x64 |
scan_response | omnibias.jax.scan_response |
soft_argmax_offset | omnibias.jax.soft_argmax_offset |
steerable_basis | omnibias.jax.steerable_basis |
tower_to_jet | omnibias.jax.tower_to_jet |
vibrational_frequencies | omnibias.jax.vibrational_frequencies |
x64_enabled | omnibias.jax.x64_enabled |