Backend switching (NumPy / JAX)¶

forge3d supports two computation backends that can be switched at runtime:

ENGINE_BACKEND=numpy python my_script.py   # default
ENGINE_BACKEND=jax   python my_script.py   # JAX JIT + vmap

Both backends produce numerically equivalent results (within float64 tolerance).

Why two backends?¶

NumPy: Simple, debuggable, no JIT overhead for small scenes.
JAX: JIT compilation + vmap batching → 2,000× throughput for RL training.

The JAX backend shines when stepping thousands of environments in parallel:

import jax
import jax.numpy as jnp
from forge3d.sim.jax_batch import batch_reach_reset, batch_reach_step

key = jax.random.PRNGKey(0)
q, tgt, obs = batch_reach_reset(key, n_envs=256)
q, obs, rew, done = batch_reach_step(q, tgt, jnp.zeros((256, 6)))

Backend API¶

backend ¶

Backend abstraction: numpy ↔ jax.

Select via environment variable before process start

ENGINE_BACKEND=numpy (default) ENGINE_BACKEND=jax

Engine code imports names from here

from forge3d.backend import xp, jit, vmap, set_at, new_prng_key, split_key, rand_uniform

xp — numpy or jax.numpy module jit — jax.jit or no-op decorator vmap — jax.vmap or loop-based fallback

backend_name ¶

backend_name() -> str

Return the active backend name ('numpy' or 'jax').

jit ¶

jit(fn: Any, **_: Any) -> Any

No-op decorator: functions run eagerly under NumPy.

vmap ¶

vmap(fn: Any, in_axes: int = 0, out_axes: int = 0) -> Any

Loop-based vmap fallback: maps fn over axis-0 of all inputs.

set_at ¶

set_at(arr: Any, idx: Any, val: Any) -> Any

Return a copy of arr with arr[idx] replaced by val.

new_prng_key ¶

new_prng_key(seed: int = 0) -> Any

Create a uint32[2] key compatible with JAX PRNG conventions.

split_key ¶

split_key(key: Any, num: int = 2) -> Any

Split key into num independent keys (deterministic, sequential seeds).

rust_core ¶

rust_core() -> Any

Rust 네이티브 확장 모듈 반환. USE_RUST_CORE=False면 None.