astronomix._pallas_helpers

Shared Pallas-backend utilities used across the FD and FV paths.

This module is the single place that: - imports Pallas / Triton (and exposes pl is None if unavailable), - normalises config.pallas_block_shape to a 3-tuple, - builds Triton CompilerParams from config knobs, - exposes the backend == PALLAS predicate, - provides the _pallas_call_sharded multi-GPU wrapper that turns an

opaque pl.pallas_call into a shard_map + ppermute halo-exchange body when the user runs on a multi-device mesh.

Every Pallas kernel module under astronomix should import from here so new knobs / fallbacks only need to be added once.

Module Contents

Functions

pallas_mesh_context	Set the active mesh for Pallas kernel sharding.
diffable_pallas_call	Run pallas_branch(state, params) with a custom_jvp boundary that routes tangent computation through native_branch.
diffable_pallas_call_n	Same as diffable_pallas_call() but takes a tuple of arbitrary differentiable primals (so callers with more than two diff args, e.g. extra rhs/accumulator buffers, can still get a custom_jvp boundary).
pallas_vjp_call	Run pallas_forward(state, aux) with a jax.custom_vjp boundary whose reverse rule is a native Pallas adjoint kernel pallas_backward.

API

astronomix._pallas_helpers.pallas_mesh_context(mesh)

Set the active mesh for Pallas kernel sharding.

time_integration enters this context around the JIT trace whenever the user supplies a sharding argument. Inside the context every Pallas kernel that calls _pallas_call_sharded will route through a shard_map + ppermute halo exchange instead of the bare pl.pallas_call.

When mesh is None (single-device run) or has size 1, the helper is a no-op — the kernel runs exactly as before.

astronomix._pallas_helpers.diffable_pallas_call(state, params, *, pallas_branch, native_branch)

Run pallas_branch(state, params) with a custom_jvp boundary that routes tangent computation through native_branch.

Both branches must accept the same positional (state, params) pair and produce the same pytree structure. Anything static (config, registered_variables, axis index, …) should be closed over.

Outside of AD the call collapses to pallas_branch(state, params) directly — no overhead. Under jax.jvp / jax.jacfwd / jax.grad / jax.vjp / jax.jacrev the custom rule fires and the tangent goes through native_branch.

astronomix._pallas_helpers.diffable_pallas_call_n(primals, *, pallas_branch, native_branch)

Same as diffable_pallas_call() but takes a tuple of arbitrary differentiable primals (so callers with more than two diff args, e.g. extra rhs/accumulator buffers, can still get a custom_jvp boundary).

astronomix._pallas_helpers.pallas_vjp_call(state, aux, *, pallas_forward, pallas_backward)

Run pallas_forward(state, aux) with a jax.custom_vjp boundary whose reverse rule is a native Pallas adjoint kernel pallas_backward.

Unlike diffable_pallas_call() (which routes the tangent — and hence the transposed gradient — through native JAX), this keeps the entire backward pass on the Pallas/GPU backend: pallas_backward(state, aux, cot) returns the input cotangent d(loss)/d(state) directly from a hand-built adjoint kernel.

Differentiates w.r.t. state only. aux (e.g. the traced SimulationParams) is threaded through the boundary and given a zero cotangent — it must be passed explicitly rather than closed over because jax.custom_vjp cannot capture traced values in its forward/backward closures (only static/concrete data — config, axis — may be closed over by the two branches). Treating the physical constants as non-differentiable matches the inverse-problem regime (gradients w.r.t. the state, not params).

NOTE: jax.custom_vjp supports reverse-mode only — jax.jvp / forward-mode AD on this boundary raises. Use it for reverse-mode (jax.grad / differentiation_mode = BACKWARDS); for forward-mode keep diffable_pallas_call().