---

name: internal-advanced description: "Use semi-public equinox.internal helpers for advanced JAX-library workflows." metadata: disco-role: operating disable-model-invocation: true license: Apache 2.0

Internal Advanced

Use this sub-skill only for advanced tasks that explicitly need equinox.internal (import equinox.internal as eqxi). This namespace is semi-public, intended for expert users and downstream JAX libraries, and has no stability guarantees.

Use when

• The user names equinox.internal, eqxi, noinline, internal while_loop, internal scan, nontraceable, nondifferentiable, closure_to_pytree, finalise_jaxpr, str2jax, GetKey, or primitive authoring helpers.
• The task is about reducing JAX compile-time by preventing inlining.
• A downstream library needs checkpointed/bounded loops or recursive checkpointed scan behavior.
• A custom primitive must support filtered PyTree arguments, JVP, transpose, or batching rules.
• A task needs advanced diagnostics that are not exposed as a stable root API.

Route elsewhere first

• Use ../module-and-trees/ for ordinary module and PyTree patterns.
• Use ../filtered-transformations/ for public filtered JIT/AD/vmap/pmap, checkpointing, callback, and sharding APIs.
• Use ../nn-and-state/ for neural-network layers and state.
• Use ../diagnostics-and-serialization/ for public debug, runtime-error, progress-meter, and serialization workflows.

Read first

• references/internal-advanced.md for supported internal families and safe usage patterns.
• references/troubleshooting.md for CPU-only, transform, loop, and primitive failure modes.
• ../../references/api-reference.md for signature summaries.

Core workflow

1. Confirm a public eqx.* API cannot solve the task. Prefer stable APIs when possible.
2. Identify the specific internal helper and its expected transform context.
3. Keep tests small and CPU-first unless the user provides accelerator-specific acceptance criteria.
4. For noinline, verify eager, jit, vmap, and AD behavior separately.
5. For internal loops, choose kind="lax", "bounded", or "checkpointed" based on gradient and memory needs.
6. For nontraceable/nondifferentiable helpers, explicitly state which transforms are meant to fail.
7. For primitive helpers, define and test impl, abstract eval, JVP, transpose, batching, and filtered binding rules.

Minimal validation

From the generated skill root (set EQUINOX_SKILL_ROOT to its absolute path), run:

cd "$EQUINOX_SKILL_ROOT" && python scripts/smoke.py --mode internal

This checks tiny noinline, internal while_loop, internal scan, and nontraceable behavior. It is a smoke test, not a guarantee that advanced JAX-internal behavior works for every backend.

Scope note

This sub-skill distills selected advanced Equinox internals into the bundled references and smoke helper. The source repository's internal modules and tests were construction evidence only, not runtime dependencies.

Key cautions

• equinox.internal APIs may change or disappear between releases.
• noinline is documented as only tested on CPU. Do not claim accelerator support without an accelerator-specific test.
• to_onnx is intentionally not a verified capability in this skill because the native ONNX test is skipped due to an upstream converter bug.
• Internal loop and primitive helpers depend on JAX internals; pin versions or re-run native tests when upgrading JAX or Equinox.