File size: 3,732 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from enum import IntEnum
from typing import Any
import equinox as eqx
import jax
import jax.numpy as jnp
from jaxtyping import Array
class MergeOpcode(IntEnum):
MERGE = 0
CARRY_LEFT = 1
CARRY_RIGHT = 2
EMPTY = 3
MERGE = int(MergeOpcode.MERGE)
CARRY_LEFT = int(MergeOpcode.CARRY_LEFT)
CARRY_RIGHT = int(MergeOpcode.CARRY_RIGHT)
EMPTY = int(MergeOpcode.EMPTY)
OPCODE_DTYPE = jnp.uint8
class QSideLeafInput(eqx.Module):
weight: Array
class FactorizedQSide(eqx.Module):
V: Array
U: Array
class SharedKernel(eqx.Module):
q_to_odd: QSideLeafInput
leaf_factors: tuple[FactorizedQSide, ...]
leaf_combiner_factors: tuple[FactorizedQSide, ...]
merge_T: Array
merge_factors: tuple[FactorizedQSide, ...]
readout_factors: tuple[FactorizedQSide, ...]
merge_eps: float = eqx.field(static=True)
class TrunkCompilerKernel(eqx.Module):
featurizer: Any
trunk: Any
shared_global: Any
class LadderProjectionKernel(eqx.Module):
weight: Array
bias: Array
norm_scale: Array
class PhysicalCompilerKernel(eqx.Module):
contextualizer: Any
global_fork: Any
leaf: Any
merge: Any
readout: Any
leaf_projection: LadderProjectionKernel
tree_pool: Any
tree_update: Any
tree_projection_weight: Array
tree_projection_bias: Array
root_projection: LadderProjectionKernel
class ModelHamiltonianArrays(eqx.Module):
coupling: Array
full_coupling: Array
field: Array
class GraphInputs(eqx.Module):
node: Array
edge: Array
class QuotientInputs(eqx.Module):
node_key: Array
edge_key: Array
class EnergyInputs(eqx.Module):
custom_lap_J_eff: Array
custom_lap_radial_const: Array
one_body_fields: tuple[Array, ...]
class EnergyMasks(eqx.Module):
real: Array
balanced: Array
class EnergyFrame(eqx.Module):
custom_lap_J_eff: Array
w_levels: tuple[Array, ...]
custom_lap_radial_const: Array
one_body_fields: tuple[Array, ...]
masks: EnergyMasks
EnergyFrameBatch = EnergyFrame
class CanonicalHamiltonian(eqx.Module):
model_coupling_fields: ModelHamiltonianArrays
node_mask: Array
balanced_mask: Array
graph_inputs: GraphInputs
quotient_inputs: QuotientInputs
energy_inputs: EnergyInputs
system_identity: Array
class SharedTrunk(eqx.Module):
node_raw: Array
edge_raw: Array
global_raw: Array
global_stream: Array
real_mask: Array
balanced_mask: Array
class CompiledTree(eqx.Module):
perm: Array
inv_perm: Array
leaf_real: Array
leaf_h: tuple[Array, ...]
leaf_combiner_h: tuple[Array, ...]
merge_h: tuple[Array, ...]
opcodes: tuple[Array, ...]
readout_h: tuple[Array, ...]
readout_combiner_h: tuple[Array, ...]
CompiledTreeBatch = CompiledTree
class CompiledWaveFunction(eqx.Module):
kernel: SharedKernel
tree: CompiledTree
def __call__(self, q_routed, _ctx=None, _t=0.0):
from .execute import execute_wavefunction
return execute_wavefunction(self.kernel, self.tree, q_routed)
class CompiledWaveFunctions(eqx.Module):
kernel: SharedKernel
trees: CompiledTreeBatch
def __call__(self, q_routed, _ctx=None, _t=0.0):
from .execute import execute_wavefunction
return execute_wavefunction(self.kernel, self.trees, q_routed)
def level_widths(n_sites: int) -> tuple[int, ...]:
if n_sites <= 0 or n_sites & (n_sites - 1):
raise ValueError(f"n_sites must be a positive power of two, got {n_sites}")
return tuple(n_sites >> level for level in range(1, n_sites.bit_length()))
|