Operators
JIT-compatible differential operators for SciML. Optimized for neural-network functions f(params, x) or f(params, t, x). Single-point (d,) or batched (N, d) inputs; batch via vmap internally.
Spatial derivatives
Operators.gradient(func, params, x)
∇f w.r.t. x
func(params, x) → scalar. x shape (d,) or (N, d). Returns gradient vector same shape as x. Pressure/temperature gradients.
Operators.jacobian(func, params, x)
Jacobian of vector field w.r.t. x
func(params, x) → vector. Returns (out_dim, in_dim) or (N, out_dim, in_dim). Velocity gradient du/dx.
Operators.divergence(func, params, x)
div(f) = trace(Jacobian)
func(params, x) → vector. Returns scalar () or (N,). Mass continuity div u.
Operators.laplacian(func, params, x)
∇²f via JVP (no full Hessian)
func(params, x) → scalar or vector. Returns Laplacian per point. Viscous terms, heat diffusion.
Operators.curl_2d(func, params, x)
Scalar curl in 2D: ∂v/∂x − ∂u/∂y
func(params, x) → [u, v]. x (2,) or (N, 2). Returns vorticity () or (N,).
Operators.curl_3d(func, params, x)
∇×f in 3D
func(params, x) → [u, v, w]. x (3,) or (N, 3). Returns (3,) or (N, 3). Vorticity; Faraday curl E.
Time derivatives
Operators.time_derivative(func, params, t, x)
∂φ/∂t
func(params, t, x) → scalar or vector. t scalar or (N,) when batched. Unsteady PDEs.
Operators.time_derivative_second(func, params, t, x)
∂²φ/∂t²
Wave equation, acoustics, second-order-in-time dynamics.
Derived
Operators.advection(u, u_grad)
(u·∇)u
u (..., d), u_grad from Operators.jacobian. Momentum equations.
Operators.symmetric_gradient(func, params, x)
Symmetric part of Jacobian (strain tensor)
func(params, x) → displacement. Returns (d, d) or (N, d, d). Hooke, strain from displacement.