File size: 2,163 Bytes
5ccb4fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import equinox as eqx
import jax
import jax.numpy as jnp
from jaxtyping import Array, Float, Int
from hamiltonzero.model.context import MultiSystemContext, SpinContext
def permute_ctx_prefix(ctx: SpinContext, perm: Int[Array, "n"]) -> SpinContext:
return eqx.tree_at(
lambda c: (
c.h_prime,
c.J_double_prime,
c.mask,
c.route_quotient_node_key,
c.route_quotient_edge_key,
),
ctx,
(
ctx.h_prime[perm],
ctx.J_double_prime[perm][:, perm, :],
ctx.mask[perm],
ctx.route_quotient_node_key[perm],
(
ctx.route_quotient_edge_key
if ctx.route_quotient_edge_key.shape[-1] == 0
else ctx.route_quotient_edge_key[perm][:, perm]
),
),
)
def permute_multi_ctx_prefix(
ctx: MultiSystemContext,
perms: Int[Array, "s n"],
) -> MultiSystemContext:
return eqx.tree_at(
lambda c: (
c.h_prime,
c.J_double_prime,
c.mask,
c.route_quotient_node_key,
c.route_quotient_edge_key,
),
ctx,
(
jax.vmap(lambda x, p: x[p])(ctx.h_prime, perms),
jax.vmap(lambda x, p: x[p][:, p, :])(ctx.J_double_prime, perms),
jax.vmap(lambda x, p: x[p])(ctx.mask, perms),
jax.vmap(lambda x, p: x[p])(ctx.route_quotient_node_key, perms),
(
ctx.route_quotient_edge_key
if ctx.route_quotient_edge_key.shape[-1] == 0
else jax.vmap(lambda x, p: x[p][:, p])(
ctx.route_quotient_edge_key, perms
)
),
),
)
def permute_q_prefix(q: Float[Array, "s b r n d"], perms: Int[Array, "s n"]):
idx = jnp.broadcast_to(perms[:, None, None, :, None], q.shape)
return jnp.take_along_axis(q, idx, axis=3)
__all__ = [
"permute_ctx_prefix",
"permute_multi_ctx_prefix",
"permute_q_prefix",
]
|