src/hamiltonzero/energy/custom_lap.py

# Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0


from __future__ import annotations

from typing import Any, Callable, NamedTuple

import jax
import jax.numpy as jnp
from jax import lax
from jax.extend import core
from jax.extend.core import Literal


class JLP(NamedTuple):
    value: Any
    jac: Any
    lap: Any
    level: int
    has_chunk_axis: bool
    chunk_idx: int = -1


def _is_jlp(x) -> bool:
    return isinstance(x, JLP)


from hamiltonzero.model._custom_lap_primitives import (
    custom_lap_active,
    enter_custom_lap,
    restore_custom_lap,
    quadrilinear_merge_p,
)


class use_custom_lap:
    def __enter__(self):
        self._cla_token = enter_custom_lap()
        return self

    def __exit__(self, *exc):
        restore_custom_lap(self._cla_token)


def build_W_levels(W_full, N: int) -> list:

    assert W_full.shape == (3 * N, 3 * N), (
        f"W_full must be [3N, 3N]; got {W_full.shape}"
    )
    assert (N & (N - 1)) == 0, f"N must be a power of 2; got {N}"
    levels = []
    k = 1
    while k <= N:
        n_chunks = N // k
        W_reshaped = W_full.reshape(n_chunks, 3 * k, n_chunks, 3 * k)
        idx = jnp.arange(n_chunks)
        W_k = W_reshaped[idx, :, idx, :]
        levels.append(W_k)
        k *= 2
    return levels


def _level_idx(k: int) -> int:

    assert k > 0 and (k & (k - 1)) == 0, f"k must be a power of 2; got {k}"
    return k.bit_length() - 1


def _W_at_level(W_levels, k: int):
    return W_levels[_level_idx(k)]


_RULE_REGISTRY: dict[core.Primitive, Callable] = {}


def _params_except(params, *drop, **defaults):

    out = {k: params[k] for k in params if k not in drop}
    for k, v in defaults.items():
        out.setdefault(k, v)
    return out


def _shape_bind_params(params, *drop):

    out = {k: params[k] for k in params if k not in drop and k != "out_sharding"}
    out.setdefault("sharding", params.get("out_sharding", None))
    return out


def _select_W_for_jlp(
    level: int, chunk_idx: int, has_chunk_axis: bool, W_levels, M_jac: int
):

    W_k = _W_at_level(W_levels, level)
    if has_chunk_axis:
        assert W_k.shape[0] == M_jac, (
            f"has_chunk_axis: W_k chunks {W_k.shape[0]} must equal jac M {M_jac}"
        )
        return W_k
    if chunk_idx >= 0:
        assert M_jac == 1
        return W_k[chunk_idx : chunk_idx + 1]

    assert W_k.shape[0] == M_jac, (
        f"multi-chunk: W_k chunks {W_k.shape[0]} must equal jac M {M_jac}"
    )
    return W_k


def _jac_self_quad_form(
    jac, level: int, chunk_idx: int, has_chunk_axis: bool, W_levels
):

    M = jac.shape[1]
    W_used = _select_W_for_jlp(level, chunk_idx, has_chunk_axis, W_levels, M)
    n_trailing = jac.ndim - 2
    if n_trailing == 0:
        out = jnp.einsum("mc,cmn,nc->c", jac, W_used, jac)
    elif n_trailing == 1:
        out = jnp.einsum("mca,cmn,nca->ca", jac, W_used, jac)
    elif n_trailing == 2:
        out = jnp.einsum("mcab,cmn,ncab->cab", jac, W_used, jac)
    elif n_trailing == 3:
        out = jnp.einsum("mcabd,cmn,ncabd->cabd", jac, W_used, jac)
    else:
        raise NotImplementedError(
            f"_jac_self_quad_form: trailing rank {n_trailing} not supported"
        )
    if not has_chunk_axis:
        out = out.sum(axis=0)
    return out


def _make_unary_rule(prim, f_prime_fn, f_dprime_fn):

    def rule(invals, params, W_levels):
        [x] = invals
        assert _is_jlp(x)
        v_out = prim.bind(x.value, **params)
        fp = f_prime_fn(x.value)
        jac_out = fp * x.jac
        fpp = f_dprime_fn(x.value)
        cross = _jac_self_quad_form(
            x.jac, x.level, x.chunk_idx, x.has_chunk_axis, W_levels
        )
        lap_out = fp * x.lap + fpp * cross
        return JLP(
            value=v_out,
            jac=jac_out,
            lap=lap_out,
            level=x.level,
            has_chunk_axis=x.has_chunk_axis,
            chunk_idx=x.chunk_idx,
        )

    _RULE_REGISTRY[prim] = rule


_make_unary_rule(lax.sin_p, jnp.cos, lambda x: -jnp.sin(x))
_make_unary_rule(lax.cos_p, lambda x: -jnp.sin(x), lambda x: -jnp.cos(x))
_make_unary_rule(
    lax.tanh_p,
    lambda x: 1.0 - jnp.tanh(x) ** 2,
    lambda x: -2.0 * jnp.tanh(x) * (1.0 - jnp.tanh(x) ** 2),
)
_make_unary_rule(lax.exp_p, jnp.exp, jnp.exp)
_make_unary_rule(lax.log_p, lambda x: 1.0 / x, lambda x: -1.0 / (x * x))
_make_unary_rule(lax.neg_p, lambda x: -jnp.ones_like(x), lambda x: jnp.zeros_like(x))
_make_unary_rule(lax.abs_p, lambda x: jnp.sign(x), lambda x: jnp.zeros_like(x))
_make_unary_rule(
    lax.sqrt_p, lambda x: 0.5 / jnp.sqrt(x), lambda x: -0.25 / (x * jnp.sqrt(x))
)
_make_unary_rule(
    lax.rsqrt_p,
    lambda x: -0.5 / (x * jnp.sqrt(x)),
    lambda x: 0.75 / (x * x * jnp.sqrt(x)),
)


def _logistic_prime(x):
    s = jax.nn.sigmoid(x)
    return s * (1.0 - s)


def _logistic_dprime(x):
    s = jax.nn.sigmoid(x)
    return s * (1.0 - s) * (1.0 - 2.0 * s)


_logistic_p = lax.logistic_p
_make_unary_rule(_logistic_p, _logistic_prime, _logistic_dprime)


def _integer_pow_rule(invals, params, W_levels):
    [x] = invals
    y = params["y"]
    assert _is_jlp(x)
    v_out = lax.integer_pow_p.bind(x.value, **params)
    if y == 0:
        return JLP(
            value=jnp.ones_like(x.value),
            jac=jnp.zeros_like(x.jac),
            lap=jnp.zeros_like(x.lap),
            level=x.level,
            has_chunk_axis=x.has_chunk_axis,
            chunk_idx=x.chunk_idx,
        )
    if y == 1:
        return x
    fp = float(y) * lax.integer_pow_p.bind(x.value, y=y - 1)
    jac_out = fp * x.jac
    fpp = float(y * (y - 1)) * lax.integer_pow_p.bind(x.value, y=max(y - 2, 0))
    cross = _jac_self_quad_form(x.jac, x.level, x.chunk_idx, x.has_chunk_axis, W_levels)
    lap_out = fp * x.lap + fpp * cross
    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=x.level,
        has_chunk_axis=x.has_chunk_axis,
        chunk_idx=x.chunk_idx,
    )


_RULE_REGISTRY[lax.integer_pow_p] = _integer_pow_rule


def _convert_rule(invals, params, W_levels):
    [x] = invals
    assert _is_jlp(x)
    new_dtype = params["new_dtype"]
    return JLP(
        value=x.value.astype(new_dtype),
        jac=x.jac.astype(new_dtype),
        lap=x.lap.astype(new_dtype),
        level=x.level,
        has_chunk_axis=x.has_chunk_axis,
        chunk_idx=x.chunk_idx,
    )


_RULE_REGISTRY[lax.convert_element_type_p] = _convert_rule


def _broadcast_jac_lap_for_op(x, target_shape):

    lap_out = jnp.broadcast_to(x.lap, target_shape)
    jac_out = _broadcast_jac_to_value_shape(x.jac, target_shape, x.has_chunk_axis)
    return jac_out, lap_out


def _broadcast_jac_to_value_shape(jac, target_value_shape, has_chunk_axis: bool):

    if has_chunk_axis:
        new_shape = (jac.shape[0], target_value_shape[0]) + tuple(
            target_value_shape[1:]
        )
    else:
        new_shape = (jac.shape[0], jac.shape[1]) + tuple(target_value_shape)
    while jac.ndim < len(new_shape):
        jac = jnp.expand_dims(jac, axis=jac.ndim)
    return jnp.broadcast_to(jac, new_shape)


def _promote_jlp_one_level(x: JLP) -> JLP:

    k = x.level
    new_k = 2 * k
    if x.has_chunk_axis:
        assert x.chunk_idx in (0, 1), (
            f"_promote_jlp_one_level (chunked): chunk_idx must be 0 or 1; got {x}"
        )
        is_left = x.chunk_idx == 0
        zero_shape = (3 * k,) + x.jac.shape[1:]
        zeros = jnp.zeros(zero_shape, dtype=x.jac.dtype)
        if is_left:
            new_jac = jnp.concatenate([x.jac, zeros], axis=0)
        else:
            new_jac = jnp.concatenate([zeros, x.jac], axis=0)
        return JLP(
            value=x.value,
            jac=new_jac,
            lap=x.lap,
            level=new_k,
            has_chunk_axis=True,
            chunk_idx=-1,
        )

    assert x.chunk_idx >= 0, (
        f"_promote_jlp_one_level: per-node requires chunk_idx>=0; got {x}"
    )
    is_left = x.chunk_idx % 2 == 0
    new_chunk_idx = x.chunk_idx // 2
    zero_shape = (3 * k,) + x.jac.shape[1:]
    zeros = jnp.zeros(zero_shape, dtype=x.jac.dtype)
    if is_left:
        new_jac = jnp.concatenate([x.jac, zeros], axis=0)
    else:
        new_jac = jnp.concatenate([zeros, x.jac], axis=0)
    return JLP(
        value=x.value,
        jac=new_jac,
        lap=x.lap,
        level=new_k,
        has_chunk_axis=False,
        chunk_idx=new_chunk_idx,
    )


def _align_jlp_levels(a: JLP, b: JLP):

    while a.level < b.level:
        a = _promote_jlp_one_level(a)
    while b.level < a.level:
        b = _promote_jlp_one_level(b)
    assert a.has_chunk_axis == b.has_chunk_axis, (
        "_align_jlp_levels: chunk-axis mismatch"
    )
    if a.has_chunk_axis:
        if (
            a.chunk_idx in (0, 1)
            and b.chunk_idx in (0, 1)
            and a.chunk_idx != b.chunk_idx
        ):
            a = _promote_jlp_one_level(a)
            b = _promote_jlp_one_level(b)
        return a, b

    while a.chunk_idx != b.chunk_idx:
        a = _promote_jlp_one_level(a)
        b = _promote_jlp_one_level(b)
    return a, b


def _add_or_sub_rule(sign: float):
    def rule(invals, params, W_levels):
        a, b = invals
        if _is_jlp(a) and _is_jlp(b):
            need_promote = (
                (a.level != b.level)
                or (
                    not a.has_chunk_axis
                    and not b.has_chunk_axis
                    and a.chunk_idx != b.chunk_idx
                )
                or (
                    a.has_chunk_axis
                    and b.has_chunk_axis
                    and a.chunk_idx in (0, 1)
                    and b.chunk_idx in (0, 1)
                    and a.chunk_idx != b.chunk_idx
                )
            )
            if need_promote:
                a, b = _align_jlp_levels(a, b)
            assert a.has_chunk_axis == b.has_chunk_axis, "add/sub: chunk-axis mismatch"
            v_out = a.value + sign * b.value
            a_jac_b = _broadcast_jac_to_value_shape(
                a.jac, v_out.shape, a.has_chunk_axis
            )
            b_jac_b = _broadcast_jac_to_value_shape(
                b.jac, v_out.shape, b.has_chunk_axis
            )
            jac_out = a_jac_b + sign * b_jac_b
            a_lap_b = jnp.broadcast_to(a.lap, v_out.shape)
            b_lap_b = jnp.broadcast_to(b.lap, v_out.shape)
            lap_out = a_lap_b + sign * b_lap_b
            return JLP(
                value=v_out,
                jac=jac_out,
                lap=lap_out,
                level=a.level,
                has_chunk_axis=a.has_chunk_axis,
                chunk_idx=a.chunk_idx,
            )
        if _is_jlp(a):
            v_out = a.value + sign * b
            jac_out, lap_out = _broadcast_jac_lap_for_op(a, v_out.shape)
            return JLP(
                value=v_out,
                jac=jac_out,
                lap=lap_out,
                level=a.level,
                has_chunk_axis=a.has_chunk_axis,
                chunk_idx=a.chunk_idx,
            )

        v_out = a + sign * b.value
        jac_out, lap_out = _broadcast_jac_lap_for_op(b, v_out.shape)
        return JLP(
            value=v_out,
            jac=sign * jac_out,
            lap=sign * lap_out,
            level=b.level,
            has_chunk_axis=b.has_chunk_axis,
            chunk_idx=b.chunk_idx,
        )

    return rule


_RULE_REGISTRY[lax.add_p] = _add_or_sub_rule(+1.0)
_RULE_REGISTRY[lax.sub_p] = _add_or_sub_rule(-1.0)


def _mul_rule(invals, params, W_levels):
    a, b = invals
    if _is_jlp(a) and _is_jlp(b):
        need_promote = (a.level != b.level) or (
            not a.has_chunk_axis and not b.has_chunk_axis and a.chunk_idx != b.chunk_idx
        )
        if need_promote:
            a, b = _align_jlp_levels(a, b)
        assert a.has_chunk_axis == b.has_chunk_axis

        v_out = a.value * b.value
        a_jac_b = _broadcast_jac_to_value_shape(a.jac, v_out.shape, a.has_chunk_axis)
        b_jac_b = _broadcast_jac_to_value_shape(b.jac, v_out.shape, b.has_chunk_axis)
        a_val_b = jnp.broadcast_to(a.value, v_out.shape)
        b_val_b = jnp.broadcast_to(b.value, v_out.shape)
        jac_out = a_val_b * b_jac_b + b_val_b * a_jac_b
        a_lap_b = jnp.broadcast_to(a.lap, v_out.shape)
        b_lap_b = jnp.broadcast_to(b.lap, v_out.shape)
        cross = _jac_cross_quad_form(
            a_jac_b, b_jac_b, a.level, a.chunk_idx, a.has_chunk_axis, W_levels
        )
        lap_out = a_val_b * b_lap_b + b_val_b * a_lap_b + 2.0 * cross
        return JLP(
            value=v_out,
            jac=jac_out,
            lap=lap_out,
            level=a.level,
            has_chunk_axis=a.has_chunk_axis,
            chunk_idx=a.chunk_idx,
        )
    if _is_jlp(a):
        v_out = a.value * b
        jac_b, lap_b = _broadcast_jac_lap_for_op(a, v_out.shape)
        return JLP(
            value=v_out,
            jac=b * jac_b,
            lap=b * lap_b,
            level=a.level,
            has_chunk_axis=a.has_chunk_axis,
            chunk_idx=a.chunk_idx,
        )

    v_out = a * b.value
    jac_b, lap_b = _broadcast_jac_lap_for_op(b, v_out.shape)
    return JLP(
        value=v_out,
        jac=a * jac_b,
        lap=a * lap_b,
        level=b.level,
        has_chunk_axis=b.has_chunk_axis,
        chunk_idx=b.chunk_idx,
    )


def _jac_cross_quad_form(
    jac_a, jac_b, level: int, chunk_idx: int, has_chunk_axis: bool, W_levels
):

    M = jac_a.shape[1]
    W_used = _select_W_for_jlp(level, chunk_idx, has_chunk_axis, W_levels, M)
    n_trailing = jac_a.ndim - 2
    if n_trailing == 0:
        out = jnp.einsum("mc,cmn,nc->c", jac_a, W_used, jac_b)
    elif n_trailing == 1:
        out = jnp.einsum("mca,cmn,nca->ca", jac_a, W_used, jac_b)
    elif n_trailing == 2:
        out = jnp.einsum("mcab,cmn,ncab->cab", jac_a, W_used, jac_b)
    elif n_trailing == 3:
        out = jnp.einsum("mcabd,cmn,ncabd->cabd", jac_a, W_used, jac_b)
    else:
        raise NotImplementedError(
            f"_jac_cross_quad_form: trailing rank {n_trailing} not supported"
        )
    if not has_chunk_axis:
        out = out.sum(axis=0)
    return out


_RULE_REGISTRY[lax.mul_p] = _mul_rule


def _div_rule(invals, params, W_levels):
    a, b = invals
    if _is_jlp(a) and _is_jlp(b):
        need_promote = (a.level != b.level) or (
            not a.has_chunk_axis and not b.has_chunk_axis and a.chunk_idx != b.chunk_idx
        )
        if need_promote:
            a, b = _align_jlp_levels(a, b)
        assert a.has_chunk_axis == b.has_chunk_axis
        v_out = a.value / b.value
        a_val_b = jnp.broadcast_to(a.value, v_out.shape)
        b_val_b = jnp.broadcast_to(b.value, v_out.shape)
        a_jac_b = _broadcast_jac_to_value_shape(a.jac, v_out.shape, a.has_chunk_axis)
        b_jac_b = _broadcast_jac_to_value_shape(b.jac, v_out.shape, b.has_chunk_axis)
        inv_b = 1.0 / b_val_b
        jac_out = inv_b * a_jac_b - (a_val_b * inv_b * inv_b) * b_jac_b
        a_lap_b = jnp.broadcast_to(a.lap, v_out.shape)
        b_lap_b = jnp.broadcast_to(b.lap, v_out.shape)
        lap_first = inv_b * a_lap_b - (a_val_b * inv_b * inv_b) * b_lap_b
        cross_bb = _jac_self_quad_form(
            b_jac_b, b.level, b.chunk_idx, b.has_chunk_axis, W_levels
        )
        cross_ab = _jac_cross_quad_form(
            a_jac_b, b_jac_b, a.level, a.chunk_idx, a.has_chunk_axis, W_levels
        )
        lap_second = (2.0 * a_val_b * inv_b**3) * cross_bb + (
            -2.0 * inv_b * inv_b
        ) * cross_ab
        lap_out = lap_first + lap_second
        return JLP(
            value=v_out,
            jac=jac_out,
            lap=lap_out,
            level=a.level,
            has_chunk_axis=a.has_chunk_axis,
            chunk_idx=a.chunk_idx,
        )
    if _is_jlp(a):
        inv_b = 1.0 / b
        v_out = a.value * inv_b
        jac_b, lap_b = _broadcast_jac_lap_for_op(a, v_out.shape)
        return JLP(
            value=v_out,
            jac=inv_b * jac_b,
            lap=inv_b * lap_b,
            level=a.level,
            has_chunk_axis=a.has_chunk_axis,
            chunk_idx=a.chunk_idx,
        )

    v_out = a / b.value
    inv_b = 1.0 / b.value
    fp = -a * inv_b * inv_b
    jac_out = fp * b.jac
    fpp = 2.0 * a * inv_b**3
    cross = _jac_self_quad_form(b.jac, b.level, b.chunk_idx, b.has_chunk_axis, W_levels)
    lap_out = fp * b.lap + fpp * cross
    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=b.level,
        has_chunk_axis=b.has_chunk_axis,
        chunk_idx=b.chunk_idx,
    )


_RULE_REGISTRY[lax.div_p] = _div_rule


def _atan2_rule(invals, params, W_levels):
    a_arg, b_arg = invals
    if _is_jlp(a_arg) and _is_jlp(b_arg):
        need_promote = (a_arg.level != b_arg.level) or (
            not a_arg.has_chunk_axis
            and not b_arg.has_chunk_axis
            and a_arg.chunk_idx != b_arg.chunk_idx
        )
        if need_promote:
            a_arg, b_arg = _align_jlp_levels(a_arg, b_arg)
        assert a_arg.has_chunk_axis == b_arg.has_chunk_axis
        a, b = a_arg, b_arg
    elif _is_jlp(a_arg):
        a = a_arg
        b = _trivial_jlp_like(b_arg, a_arg)
    else:
        b = b_arg
        a = _trivial_jlp_like(a_arg, b_arg)
    v_out = jnp.arctan2(a.value, b.value)
    r_sq = a.value**2 + b.value**2
    inv_r2 = 1.0 / r_sq
    fa = b.value * inv_r2
    fb = -a.value * inv_r2
    aj = _broadcast_jac_to_value_shape(a.jac, v_out.shape, a.has_chunk_axis)
    bj = _broadcast_jac_to_value_shape(b.jac, v_out.shape, b.has_chunk_axis)
    jac_out = fa * aj + fb * bj
    faa = -2.0 * a.value * b.value * inv_r2 * inv_r2
    fbb = 2.0 * a.value * b.value * inv_r2 * inv_r2
    fab = (a.value**2 - b.value**2) * inv_r2 * inv_r2
    al = jnp.broadcast_to(a.lap, v_out.shape)
    bl = jnp.broadcast_to(b.lap, v_out.shape)
    cross_aa = _jac_self_quad_form(aj, a.level, a.chunk_idx, a.has_chunk_axis, W_levels)
    cross_bb = _jac_self_quad_form(bj, b.level, b.chunk_idx, b.has_chunk_axis, W_levels)
    cross_ab = _jac_cross_quad_form(
        aj, bj, a.level, a.chunk_idx, a.has_chunk_axis, W_levels
    )
    lap_out = fa * al + fb * bl + faa * cross_aa + fbb * cross_bb + 2.0 * fab * cross_ab
    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=a.level,
        has_chunk_axis=a.has_chunk_axis,
        chunk_idx=a.chunk_idx,
    )


def _trivial_jlp_like(plain_val, template_jlp: JLP) -> JLP:

    val = jnp.asarray(plain_val)
    if template_jlp.has_chunk_axis:
        jac_shape = (template_jlp.jac.shape[0], template_jlp.jac.shape[1]) + tuple(
            val.shape[1:]
        )
    else:
        jac_shape = (template_jlp.jac.shape[0], template_jlp.jac.shape[1]) + tuple(
            val.shape
        )
    return JLP(
        value=val,
        jac=jnp.zeros(jac_shape, dtype=val.dtype),
        lap=jnp.zeros_like(val),
        level=template_jlp.level,
        has_chunk_axis=template_jlp.has_chunk_axis,
        chunk_idx=template_jlp.chunk_idx,
    )


_RULE_REGISTRY[lax.atan2_p] = _atan2_rule


def _dot_general_rule(invals, params, W_levels):
    lhs, rhs = invals
    dimension_numbers = params["dimension_numbers"]
    (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dimension_numbers

    if _is_jlp(lhs) and _is_jlp(rhs):
        raise NotImplementedError(
            "dot_general: JLP × JLP outside quadrilinear_merge_p is not supported. "
            "If a model op needs same-level bilinear, route it through the merge "
            "primitive or rewrite as elementwise mul + reduce_sum."
        )
    if _is_jlp(lhs):
        return _dot_general_jlp_plain(
            lhs, rhs, dimension_numbers, params, jlp_is_lhs=True
        )
    return _dot_general_jlp_plain(rhs, lhs, dimension_numbers, params, jlp_is_lhs=False)


def _dot_general_jlp_plain(
    jlp_arg, plain_arg, dimension_numbers, params, jlp_is_lhs: bool
):

    (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dimension_numbers
    dot_kw = {
        "precision": params.get("precision", None),
        "preferred_element_type": params.get("preferred_element_type", None),
        "out_sharding": params.get("out_sharding", None),
    }

    if jlp_is_lhs:
        v_out = lax.dot_general(jlp_arg.value, plain_arg, dimension_numbers, **dot_kw)
        lap_out = lax.dot_general(jlp_arg.lap, plain_arg, dimension_numbers, **dot_kw)
    else:
        v_out = lax.dot_general(plain_arg, jlp_arg.value, dimension_numbers, **dot_kw)
        lap_out = lax.dot_general(plain_arg, jlp_arg.lap, dimension_numbers, **dot_kw)

    jac = jlp_arg.jac
    leading_3k = jac.shape[0]

    if jlp_arg.has_chunk_axis:
        jac_for_dot = jac
        shift = 1
    else:
        jac_for_dot = jac.reshape((leading_3k,) + tuple(jlp_arg.value.shape))
        shift = 1

    if jlp_is_lhs:
        new_lhs_contract = tuple(a + shift for a in lhs_contract)
        new_rhs_contract = tuple(rhs_contract)
        new_lhs_batch = tuple(a + shift for a in lhs_batch)
        new_rhs_batch = tuple(rhs_batch)
        new_dim_nums = (
            (new_lhs_contract, new_rhs_contract),
            (new_lhs_batch, new_rhs_batch),
        )
        jac_out_raw = lax.dot_general(jac_for_dot, plain_arg, new_dim_nums, **dot_kw)

        n_batch = len(new_lhs_batch)
        pos_3k = n_batch
    else:
        new_lhs_contract = tuple(lhs_contract)
        new_rhs_contract = tuple(a + shift for a in rhs_contract)
        new_lhs_batch = tuple(lhs_batch)
        new_rhs_batch = tuple(a + shift for a in rhs_batch)
        new_dim_nums = (
            (new_lhs_contract, new_rhs_contract),
            (new_lhs_batch, new_rhs_batch),
        )
        jac_out_raw = lax.dot_general(plain_arg, jac_for_dot, new_dim_nums, **dot_kw)

        n_batch = len(new_lhs_batch)
        lhs_ndim = jnp.asarray(plain_arg).ndim
        n_lhs_nonbatch = lhs_ndim - n_batch - len(new_lhs_contract)
        pos_3k = n_batch + n_lhs_nonbatch

    if pos_3k != 0:
        jac_out = jnp.moveaxis(jac_out_raw, pos_3k, 0)
    else:
        jac_out = jac_out_raw

    if jlp_arg.has_chunk_axis:
        new_has_chunk_axis = True
        new_chunk_idx = -1
    else:
        jac_out = jac_out[:, None]
        new_has_chunk_axis = False
        new_chunk_idx = jlp_arg.chunk_idx

    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=jlp_arg.level,
        has_chunk_axis=new_has_chunk_axis,
        chunk_idx=new_chunk_idx,
    )


_RULE_REGISTRY[lax.dot_general_p] = _dot_general_rule


def _broadcast_in_dim_rule(invals, params, W_levels):
    [x] = invals
    assert _is_jlp(x)
    shape = params["shape"]
    broadcast_dimensions = params["broadcast_dimensions"]
    extra = _shape_bind_params(params, "shape", "broadcast_dimensions")

    v_out = lax.broadcast_in_dim_p.bind(
        x.value, shape=shape, broadcast_dimensions=broadcast_dimensions, **extra
    )
    lap_out = lax.broadcast_in_dim_p.bind(
        x.lap, shape=shape, broadcast_dimensions=broadcast_dimensions, **extra
    )

    if x.has_chunk_axis:
        new_shape = (x.jac.shape[0],) + tuple(shape)
        new_bd = (0,) + tuple(d + 1 for d in broadcast_dimensions)
        jac_out = lax.broadcast_in_dim_p.bind(
            x.jac.reshape((x.jac.shape[0],) + x.value.shape),
            shape=new_shape,
            broadcast_dimensions=new_bd,
            **extra,
        )

        new_has_chunk_axis = True
        new_chunk_idx = -1
    else:
        new_shape = (x.jac.shape[0], 1) + tuple(shape)
        new_bd = (0, 1) + tuple(d + 2 for d in broadcast_dimensions)
        jac_out = lax.broadcast_in_dim_p.bind(
            x.jac,
            shape=new_shape,
            broadcast_dimensions=new_bd,
            **extra,
        )
        new_has_chunk_axis = False
        new_chunk_idx = x.chunk_idx

    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=x.level,
        has_chunk_axis=new_has_chunk_axis,
        chunk_idx=new_chunk_idx,
    )


_RULE_REGISTRY[lax.broadcast_in_dim_p] = _broadcast_in_dim_rule


def _reduce_sum_rule(invals, params, W_levels):
    [x] = invals
    assert _is_jlp(x)
    axes = tuple(params["axes"])
    extra = _params_except(params, "axes", out_sharding=None)
    v_out = lax.reduce_sum_p.bind(x.value, axes=axes, **extra)
    lap_out = lax.reduce_sum_p.bind(x.lap, axes=axes, **extra)
    if x.has_chunk_axis:
        if 0 in axes:
            non_chunk_axes = tuple(a for a in axes if a != 0)
            jac_reduce_axes = tuple(a + 1 for a in non_chunk_axes)
            if jac_reduce_axes:
                jac_out = lax.reduce_sum_p.bind(x.jac, axes=jac_reduce_axes, **extra)
            else:
                jac_out = x.jac
            new_has_chunk_axis = False
            new_chunk_idx = -1
        else:
            jac_reduce_axes = tuple(a + 1 for a in axes)
            jac_out = lax.reduce_sum_p.bind(x.jac, axes=jac_reduce_axes, **extra)
            new_has_chunk_axis = True
            new_chunk_idx = -1
    else:
        jac_reduce_axes = tuple(a + 2 for a in axes)
        jac_out = lax.reduce_sum_p.bind(x.jac, axes=jac_reduce_axes, **extra)
        new_has_chunk_axis = False
        new_chunk_idx = x.chunk_idx
    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=x.level,
        has_chunk_axis=new_has_chunk_axis,
        chunk_idx=new_chunk_idx,
    )


_RULE_REGISTRY[lax.reduce_sum_p] = _reduce_sum_rule


def _reduce_max_rule(invals, params, W_levels):
    [x] = invals
    assert _is_jlp(x)
    axes = tuple(params["axes"])
    extra = _params_except(params, "axes", out_sharding=None)
    v_out = lax.reduce_max_p.bind(x.value, axes=axes, **extra)

    keep_shape = list(x.value.shape)
    for a in axes:
        keep_shape[a] = 1
    v_max_kept = lax.reduce_max_p.bind(x.value, axes=axes, **extra).reshape(
        tuple(keep_shape)
    )
    mask = (x.value == v_max_kept).astype(x.value.dtype)

    mask_sum_axes = lax.reduce_sum_p.bind(
        mask,
        axes=axes,
        **extra,
    ).reshape(tuple(keep_shape))
    mask = mask / (mask_sum_axes + 1e-30)

    def _gated_reduce(jac_arr, jac_axes):

        n_lead = jac_arr.ndim - x.value.ndim
        mask_b = mask.reshape((1,) * n_lead + mask.shape)

        return lax.reduce_sum_p.bind(
            jac_arr * mask_b,
            axes=jac_axes,
            **extra,
        )

    if x.has_chunk_axis:
        if 0 in axes:
            non_chunk_axes = tuple(a for a in axes if a != 0)
            jac_reduce_axes = tuple(a + 1 for a in non_chunk_axes) + (1,)
            jac_out = _gated_reduce(x.jac, jac_reduce_axes)
            new_has_chunk_axis = False
            new_chunk_idx = -1
        else:
            jac_reduce_axes = tuple(a + 1 for a in axes)
            jac_out = _gated_reduce(x.jac, jac_reduce_axes)
            new_has_chunk_axis = True
            new_chunk_idx = -1
    else:
        jac_reduce_axes = tuple(a + 2 for a in axes)
        jac_out = _gated_reduce(x.jac, jac_reduce_axes)
        new_has_chunk_axis = False
        new_chunk_idx = x.chunk_idx
    lap_out = jnp.zeros_like(v_out)
    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=x.level,
        has_chunk_axis=new_has_chunk_axis,
        chunk_idx=new_chunk_idx,
    )


_RULE_REGISTRY[lax.reduce_max_p] = _reduce_max_rule


def _make_minmax_rule(sign_for_first):

    def rule(invals, params, W_levels):
        a, b = invals
        if _is_jlp(a) and _is_jlp(b):
            need_promote = (
                (a.level != b.level)
                or (
                    not a.has_chunk_axis
                    and not b.has_chunk_axis
                    and a.chunk_idx != b.chunk_idx
                )
                or (
                    a.has_chunk_axis
                    and b.has_chunk_axis
                    and a.chunk_idx in (0, 1)
                    and b.chunk_idx in (0, 1)
                    and a.chunk_idx != b.chunk_idx
                )
            )
            if need_promote:
                a, b = _align_jlp_levels(a, b)
            cmp = (a.value - b.value) * sign_for_first
            mask_a = (cmp > 0).astype(a.value.dtype)
            mask_b = 1.0 - mask_a
            v_out = mask_a * a.value + mask_b * b.value

            jac_a_b = _broadcast_jac_to_value_shape(
                a.jac, v_out.shape, a.has_chunk_axis
            )
            jac_b_b = _broadcast_jac_to_value_shape(
                b.jac, v_out.shape, b.has_chunk_axis
            )
            n_lead_a = jac_a_b.ndim - mask_a.ndim
            n_lead_b = jac_b_b.ndim - mask_b.ndim
            mask_a_lead = mask_a.reshape((1,) * n_lead_a + mask_a.shape)
            mask_b_lead = mask_b.reshape((1,) * n_lead_b + mask_b.shape)
            jac_out = mask_a_lead * jac_a_b + mask_b_lead * jac_b_b
            lap_out = mask_a * jnp.broadcast_to(
                a.lap, v_out.shape
            ) + mask_b * jnp.broadcast_to(b.lap, v_out.shape)
            return JLP(
                value=v_out,
                jac=jac_out,
                lap=lap_out,
                level=a.level,
                has_chunk_axis=a.has_chunk_axis,
                chunk_idx=a.chunk_idx,
            )
        if _is_jlp(a):
            cmp = (a.value - b) * sign_for_first
            mask_a = (cmp > 0).astype(a.value.dtype)
            v_out = mask_a * a.value + (1 - mask_a) * b
            n_lead = a.jac.ndim - mask_a.ndim
            mask_a_lead = mask_a.reshape((1,) * n_lead + mask_a.shape)
            jac_out = mask_a_lead * a.jac
            lap_out = mask_a * a.lap
            return JLP(
                value=v_out,
                jac=jac_out,
                lap=lap_out,
                level=a.level,
                has_chunk_axis=a.has_chunk_axis,
                chunk_idx=a.chunk_idx,
            )

        cmp = (a - b.value) * sign_for_first
        mask_a = (cmp > 0).astype(b.value.dtype)
        v_out = mask_a * a + (1 - mask_a) * b.value
        mask_b = 1 - mask_a
        n_lead = b.jac.ndim - mask_b.ndim
        mask_b_lead = mask_b.reshape((1,) * n_lead + mask_b.shape)
        jac_out = mask_b_lead * b.jac
        lap_out = mask_b * b.lap
        return JLP(
            value=v_out,
            jac=jac_out,
            lap=lap_out,
            level=b.level,
            has_chunk_axis=b.has_chunk_axis,
            chunk_idx=b.chunk_idx,
        )

    return rule


_RULE_REGISTRY[lax.max_p] = _make_minmax_rule(+1.0)
_RULE_REGISTRY[lax.min_p] = _make_minmax_rule(-1.0)


def _reshape_jac_for_value(x: JLP, new_value_shape: tuple, new_dimensions):

    leading_3k = x.jac.shape[0]
    if x.has_chunk_axis:
        M = x.value.shape[0]
        old_trailing = x.value.shape[1:]

        if len(new_value_shape) >= 1 and new_value_shape[0] == M:
            jac_axes_perm = None
            if new_dimensions is not None:
                jac_axes_perm = (0,) + tuple(d + 1 for d in new_dimensions)
                jac_pre = jnp.transpose(x.jac, jac_axes_perm)
            else:
                jac_pre = x.jac
            new_jac_shape = (leading_3k,) + tuple(new_value_shape)
            jac_new = jnp.reshape(jac_pre, new_jac_shape)

            return jac_new, True, -1

        if M == 1:
            assert new_dimensions is None, (
                "reshape with M=1-squeeze + transpose not yet supported"
            )

            new_jac_shape = (leading_3k, 1) + tuple(new_value_shape)

            jac_new = jnp.reshape(x.jac, new_jac_shape)
            return jac_new, False, 0
        raise NotImplementedError(
            f"reshape: cannot reshape JLP value {x.value.shape} (has_chunk_axis, M={M}) "
            f"to {new_value_shape} — chunk axis would be fused/lost."
        )

    assert new_dimensions is None or all(d >= 0 for d in new_dimensions)
    jac_pre = x.jac
    if new_dimensions is not None:
        jac_axes_perm = (0, 1) + tuple(d + 2 for d in new_dimensions)
        jac_pre = jnp.transpose(jac_pre, jac_axes_perm)
    new_jac_shape = (leading_3k, 1) + tuple(new_value_shape)
    jac_new = jnp.reshape(jac_pre, new_jac_shape)
    return jac_new, False, x.chunk_idx


def _reshape_rule(invals, params, W_levels):
    [x] = invals
    assert _is_jlp(x)
    new_sizes = tuple(params["new_sizes"])
    dimensions = params.get("dimensions")

    extra = _shape_bind_params(params, "new_sizes", "dimensions")
    v_out = lax.reshape_p.bind(
        x.value, new_sizes=new_sizes, dimensions=dimensions, **extra
    )
    lap_out = lax.reshape_p.bind(
        x.lap, new_sizes=new_sizes, dimensions=dimensions, **extra
    )
    jac_new, new_has_chunk, new_idx = _reshape_jac_for_value(x, new_sizes, dimensions)
    return JLP(
        value=v_out,
        jac=jac_new,
        lap=lap_out,
        level=x.level,
        has_chunk_axis=new_has_chunk,
        chunk_idx=new_idx,
    )


_RULE_REGISTRY[lax.reshape_p] = _reshape_rule


def _transpose_rule(invals, params, W_levels):
    [x] = invals
    assert _is_jlp(x)
    perm = tuple(params["permutation"])
    extra = {k: params[k] for k in params if k != "permutation"}
    v_out = lax.transpose_p.bind(x.value, permutation=perm, **extra)
    lap_out = lax.transpose_p.bind(x.lap, permutation=perm, **extra)

    if x.has_chunk_axis:
        if perm[0] != 0:
            raise NotImplementedError(
                "transpose: chunk axis (value axis 0) must remain at position 0; "
                f"got permutation {perm}."
            )

        jac_perm = (0, 1) + tuple(p + 1 for p in perm[1:])
    else:
        jac_perm = (0, 1) + tuple(p + 2 for p in perm)
    jac_new = lax.transpose_p.bind(x.jac, permutation=jac_perm, **extra)
    return JLP(
        value=v_out,
        jac=jac_new,
        lap=lap_out,
        level=x.level,
        has_chunk_axis=x.has_chunk_axis,
        chunk_idx=x.chunk_idx,
    )


_RULE_REGISTRY[lax.transpose_p] = _transpose_rule


def _slice_rule(invals, params, W_levels):
    [x] = invals
    assert _is_jlp(x)
    start = tuple(params["start_indices"])
    limit = tuple(params["limit_indices"])
    strides = params.get("strides")
    extra = {
        k: params[k]
        for k in params
        if k not in ("start_indices", "limit_indices", "strides")
    }
    v_out = lax.slice_p.bind(
        x.value, start_indices=start, limit_indices=limit, strides=strides, **extra
    )
    lap_out = lax.slice_p.bind(
        x.lap, start_indices=start, limit_indices=limit, strides=strides, **extra
    )

    if x.has_chunk_axis:
        new_chunk_idx = x.chunk_idx
        old_M = x.value.shape[0]
        stride0 = strides[0] if strides is not None else 1
        v_out_M = v_out.shape[0]
        chunk_axis_touched = start[0] != 0 or limit[0] != old_M or stride0 != 1
        if chunk_axis_touched:
            if v_out_M == 1:
                new_chunk_idx = start[0]
            elif stride0 > 1 and v_out_M * stride0 == old_M and start[0] in (0, 1):
                new_chunk_idx = start[0]
            elif v_out_M != old_M:
                raise NotImplementedError(
                    f"slice on chunk axis: unsupported sub-range "
                    f"start={start[0]}, limit={limit[0]}, stride={stride0}, "
                    f"old_M={old_M}, v_out_M={v_out_M}"
                )
        jac_start = (0, start[0]) + tuple(start[1:])
        jac_limit = (x.jac.shape[0], limit[0]) + tuple(limit[1:])
        jac_strides = None if strides is None else (1, strides[0]) + tuple(strides[1:])
        jac_out = lax.slice_p.bind(
            x.jac,
            start_indices=jac_start,
            limit_indices=jac_limit,
            strides=jac_strides,
            **extra,
        )
        return JLP(
            value=v_out,
            jac=jac_out,
            lap=lap_out,
            level=x.level,
            has_chunk_axis=True,
            chunk_idx=new_chunk_idx,
        )

    jac_start = (0, 0) + tuple(start)
    jac_limit = (x.jac.shape[0], x.jac.shape[1]) + tuple(limit)
    jac_strides = None if strides is None else (1, 1) + tuple(strides)
    jac_out = lax.slice_p.bind(
        x.jac,
        start_indices=jac_start,
        limit_indices=jac_limit,
        strides=jac_strides,
        **extra,
    )
    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=x.level,
        has_chunk_axis=False,
        chunk_idx=x.chunk_idx,
    )


_RULE_REGISTRY[lax.slice_p] = _slice_rule


def _squeeze_rule(invals, params, W_levels):
    [x] = invals
    assert _is_jlp(x)
    dims = tuple(params["dimensions"])
    extra = {k: params[k] for k in params if k != "dimensions"}
    v_out = lax.squeeze_p.bind(x.value, dimensions=dims, **extra)
    lap_out = lax.squeeze_p.bind(x.lap, dimensions=dims, **extra)
    if x.has_chunk_axis and 0 in dims:
        non_chunk_dims = tuple(d for d in dims if d != 0)
        jac_dims = tuple(d + 1 for d in non_chunk_dims)

        if jac_dims:
            jac_out = lax.squeeze_p.bind(x.jac, dimensions=jac_dims, **extra)
        else:
            jac_out = x.jac
        return JLP(
            value=v_out,
            jac=jac_out,
            lap=lap_out,
            level=x.level,
            has_chunk_axis=False,
            chunk_idx=x.chunk_idx,
        )
    if x.has_chunk_axis:
        jac_dims = tuple(d + 1 for d in dims)
    else:
        jac_dims = tuple(d + 2 for d in dims)
    if jac_dims:
        jac_out = lax.squeeze_p.bind(x.jac, dimensions=jac_dims, **extra)
    else:
        jac_out = x.jac
    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=x.level,
        has_chunk_axis=x.has_chunk_axis,
        chunk_idx=x.chunk_idx,
    )


_RULE_REGISTRY[lax.squeeze_p] = _squeeze_rule


_stack_p = lax.stack_p


def _stack_rule(invals, params, W_levels):
    del W_levels
    axis = int(params["axis"])
    extra = {k: params[k] for k in params if k != "axis"}

    jlp_inputs = [v for v in invals if _is_jlp(v)]
    levels = {v.level for v in jlp_inputs}
    assert len(levels) == 1, f"stack: mixed JLP levels {levels}"
    level = next(iter(levels))
    has_chunk = jlp_inputs[0].has_chunk_axis
    chunk_idx = jlp_inputs[0].chunk_idx
    for v in jlp_inputs[1:]:
        assert v.has_chunk_axis == has_chunk and v.chunk_idx == chunk_idx, (
            "stack: inconsistent chunk metadata"
        )

    if has_chunk and axis == 0:
        raise NotImplementedError(
            "stack: inserting an axis before the JLP chunk axis is not supported"
        )

    values = [v.value if _is_jlp(v) else v for v in invals]
    value_out = _stack_p.bind(*values, axis=axis, **extra)

    laps = [v.lap if _is_jlp(v) else jnp.zeros_like(v) for v in invals]
    lap_out = _stack_p.bind(*laps, axis=axis, **extra)

    jac_axis = axis + (1 if has_chunk else 2)
    ref_jac = jlp_inputs[0].jac
    jacs = [v.jac if _is_jlp(v) else jnp.zeros_like(ref_jac) for v in invals]
    jac_out = _stack_p.bind(*jacs, axis=jac_axis, **extra)
    return JLP(
        value=value_out,
        jac=jac_out,
        lap=lap_out,
        level=level,
        has_chunk_axis=has_chunk,
        chunk_idx=chunk_idx,
    )


_RULE_REGISTRY[_stack_p] = _stack_rule


def _concatenate_rule(invals, params, W_levels):
    dim = params["dimension"]
    extra = {k: params[k] for k in params if k != "dimension"}

    jlp_inputs = [v for v in invals if _is_jlp(v)]
    levels = set(v.level for v in jlp_inputs)
    assert len(levels) == 1, f"concatenate: mixed JLP levels {levels}"
    level = next(iter(levels))
    has_chunk = jlp_inputs[0].has_chunk_axis
    chunk_idx = jlp_inputs[0].chunk_idx
    for v in jlp_inputs[1:]:
        assert v.has_chunk_axis == has_chunk and v.chunk_idx == chunk_idx, (
            "concatenate: inconsistent chunk metadata"
        )
    if has_chunk and dim == 0:
        raise NotImplementedError("concatenate on chunk axis not supported")

    values = [v.value if _is_jlp(v) else v for v in invals]
    v_out = lax.concatenate_p.bind(*values, dimension=dim, **extra)
    laps = [v.lap if _is_jlp(v) else jnp.zeros_like(v) for v in invals]
    lap_out = lax.concatenate_p.bind(*laps, dimension=dim, **extra)

    jac_dim = dim + 2 if not has_chunk else dim + 1
    jacs = []
    for v in invals:
        if _is_jlp(v):
            jacs.append(v.jac)
        else:
            shape_v = v.shape if hasattr(v, "shape") else jnp.asarray(v).shape
            if has_chunk:
                M = jlp_inputs[0].jac.shape[1]
                jac_shape = (jlp_inputs[0].jac.shape[0], M) + tuple(shape_v[1:])
            else:
                M = jlp_inputs[0].jac.shape[1]
                jac_shape = (jlp_inputs[0].jac.shape[0], M) + tuple(shape_v)
            jacs.append(jnp.zeros(jac_shape, dtype=jlp_inputs[0].jac.dtype))
    jac_out = lax.concatenate_p.bind(*jacs, dimension=jac_dim, **extra)
    return JLP(
        value=v_out,
        jac=jac_out,
        lap=lap_out,
        level=level,
        has_chunk_axis=has_chunk,
        chunk_idx=chunk_idx,
    )


_RULE_REGISTRY[lax.concatenate_p] = _concatenate_rule


def _jit_p_rule(invals, params, W_levels):

    inner_jaxpr = params["jaxpr"]
    j = inner_jaxpr.jaxpr
    consts = inner_jaxpr.consts
    env: dict = {}
    for cv, c in zip(j.constvars, consts):
        env[cv] = c
    for iv, x in zip(j.invars, invals):
        env[iv] = x
    for eqn in j.eqns:
        outvals = _eval_eqn(eqn, env, W_levels)
        for ov, ov_val in zip(eqn.outvars, outvals):
            env[ov] = ov_val
    return [env[ov] for ov in j.outvars]


from jax._src import pjit as _pjit_module

_RULE_REGISTRY[_pjit_module.jit_p] = _jit_p_rule


def _quadrilinear_merge_rule(invals, params, W_levels):

    T, u_a, u_b = invals
    assert not _is_jlp(T), "quadrilinear_merge: T must be plain"
    assert _is_jlp(u_a) and _is_jlp(u_b), "quadrilinear_merge: u_a, u_b must be JLPs"
    assert u_a.level == u_b.level, (
        f"quadrilinear_merge: leg levels differ {u_a.level} vs {u_b.level}"
    )
    assert u_a.has_chunk_axis and u_b.has_chunk_axis, (
        "quadrilinear_merge requires the compiled chunk axis"
    )

    k = u_a.level
    new_k = 2 * k
    G, d_r, _, _ = T.shape
    d_m_eff = G * d_r
    M = u_a.value.shape[0]
    assert u_b.value.shape[0] == M, "quadrilinear_merge: chunked legs must agree on M"
    u_a_2d = u_a.value.reshape(M, G, d_r)
    u_b_2d = u_b.value.reshape(M, G, d_r)
    raw_value_2d = jnp.einsum("ijkl,mik,mil->mij", T, u_a_2d, u_b_2d)
    raw_value = raw_value_2d.reshape(M, d_m_eff)
    ua_jac_2d = u_a.jac.reshape(u_a.jac.shape[0], M, G, d_r)
    ub_jac_2d = u_b.jac.reshape(u_b.jac.shape[0], M, G, d_r)
    jac_upper_2d = jnp.einsum(
        "ijkl,Amik,mil->Amij",
        T,
        ua_jac_2d,
        u_b_2d,
    )
    jac_lower_2d = jnp.einsum(
        "ijkl,mik,Bmil->Bmij",
        T,
        u_a_2d,
        ub_jac_2d,
    )
    jac_upper = jac_upper_2d.reshape(jac_upper_2d.shape[0], M, d_m_eff)
    jac_lower = jac_lower_2d.reshape(jac_lower_2d.shape[0], M, d_m_eff)
    jac_out = jnp.concatenate([jac_upper, jac_lower], axis=0)
    ua_lap_2d = u_a.lap.reshape(M, G, d_r)
    ub_lap_2d = u_b.lap.reshape(M, G, d_r)
    term1_2d = jnp.einsum(
        "ijkl,mik,mil->mij",
        T,
        ua_lap_2d,
        u_b_2d,
    )
    term2_2d = jnp.einsum(
        "ijkl,mik,mil->mij",
        T,
        u_a_2d,
        ub_lap_2d,
    )
    W_2k = _W_at_level(W_levels, new_k)
    W_off = W_2k[:, : 3 * k, 3 * k :]
    cross_2d = jnp.einsum(
        "ijkl,Amik,Bmil,mAB->mij",
        T,
        ua_jac_2d,
        ub_jac_2d,
        W_off,
    )
    lap_out_2d = term1_2d + term2_2d + 2.0 * cross_2d
    lap_out = lap_out_2d.reshape(M, d_m_eff)
    return JLP(
        value=raw_value,
        jac=jac_out,
        lap=lap_out,
        level=new_k,
        has_chunk_axis=True,
        chunk_idx=-1,
    )


_RULE_REGISTRY[quadrilinear_merge_p] = _quadrilinear_merge_rule


def _eval_eqn(eqn, env, W_levels):
    invals = []
    for v in eqn.invars:
        if isinstance(v, Literal):
            invals.append(v.val)
        else:
            invals.append(env[v])
    has_jlp = any(_is_jlp(x) for x in invals)
    if not has_jlp:
        bind_params = eqn.primitive.get_bind_params(eqn.params)
        outvals = eqn.primitive.bind(*invals, **bind_params)
        if not eqn.primitive.multiple_results:
            outvals = [outvals]
        return outvals
    rule = _RULE_REGISTRY.get(eqn.primitive)
    if rule is None:
        raise NotImplementedError(
            f"custom_lap: no rule for primitive {eqn.primitive.name!r}. "
            f"eqn = {eqn}. Register a rule in energy/custom_lap.py."
        )
    outvals = rule(invals, eqn.params, W_levels)
    if not eqn.primitive.multiple_results:
        outvals = [outvals]
    return outvals


def _trace_z(fn, z, N, W_levels):

    z = jnp.asarray(z)
    assert z.shape == (N, 3), f"z must be [N={N}, 3]; got {z.shape}"
    eye3 = jnp.eye(3, dtype=z.dtype)
    z_jac = jnp.broadcast_to(eye3[:, None, :], (3, N, 3))
    z_lap = jnp.zeros((N, 3), dtype=z.dtype)
    z_jlp = JLP(
        value=z,
        jac=z_jac,
        lap=z_lap,
        level=1,
        has_chunk_axis=True,
        chunk_idx=-1,
    )
    closed = jax.make_jaxpr(fn)(z)
    jaxpr = closed.jaxpr
    consts = closed.consts
    env: dict = {}
    for cv, c in zip(jaxpr.constvars, consts):
        env[cv] = c
    env[jaxpr.invars[0]] = z_jlp
    for eqn in jaxpr.eqns:
        outvals = _eval_eqn(eqn, env, W_levels)
        for ov, ov_val in zip(eqn.outvars, outvals):
            env[ov] = ov_val
    return [env[ov] for ov in jaxpr.outvars]


def _canonical_jac(jlp: JLP):

    if jlp.jac.shape[1] != 1:
        jac = jnp.swapaxes(jlp.jac, 0, 1)
        return jac.reshape((jlp.jac.shape[0] * jlp.jac.shape[1],) + jlp.jac.shape[2:])
    return jlp.jac[:, 0]


def custom_forward_laplacian_with_jac(fn: Callable, W_levels: list, N: int) -> Callable:

    def lap_jac_fn(z):
        outs = _trace_z(fn, z, N, W_levels)
        if len(outs) == 1:
            out = outs[0]
            if _is_jlp(out):
                return out.value, _canonical_jac(out), out.lap
            value = out
            return (
                value,
                jnp.zeros((3 * N,) + value.shape, dtype=value.dtype),
                jnp.zeros_like(value),
            )
        values = tuple(o.value if _is_jlp(o) else o for o in outs)
        jacs = tuple(
            _canonical_jac(o)
            if _is_jlp(o)
            else jnp.zeros((3 * N,) + o.shape, dtype=o.dtype)
            for o in outs
        )
        laps = tuple(o.lap if _is_jlp(o) else jnp.zeros_like(o) for o in outs)
        return values, jacs, laps

    return lap_jac_fn


__all__ = [
    "JLP",
    "build_W_levels",
    "custom_forward_laplacian_with_jac",
    "custom_lap_active",
    "use_custom_lap",
]