| |
| |
|
|
| from __future__ import annotations |
| import equinox as eqx |
| import jax |
| import jax.numpy as jnp |
| from jaxtyping import Array, Float, Int, PRNGKeyArray |
| from hamiltonzero.model.odd_ops import BiasFreeLinear, Linear, _RMS |
| from hamiltonzero.model.fused_silu import fused_silu |
|
|
|
|
| class GroupLayerNorm(eqx.Module): |
| weight: Float[Array, "n_groups d_group"] |
| eps: float = eqx.field(static=True, default=1e-05) |
| _use_id: str = eqx.field(static=True, default="") |
| n_groups: int = eqx.field(static=True) |
| d_group: int = eqx.field(static=True) |
|
|
| def __init__(self, n_groups: int, d_group: int, eps: float = 1e-05): |
| self.weight = jnp.ones((n_groups, d_group)) |
| self.eps = float(eps) |
| self._use_id = "" |
| self.n_groups = int(n_groups) |
| self.d_group = int(d_group) |
|
|
| def __call__( |
| self, |
| x: Float[Array, "... n_groups d_group"], |
| *, |
| pathway: str | None = None, |
| kfac_structural_mask=None, |
| kfac_repeat_ndim: int = 0, |
| kfac_context_primal_reused_over_walkers: bool = False, |
| ) -> Float[Array, "... n_groups d_group"]: |
| if x.shape[-2] != self.n_groups or x.shape[-1] != self.d_group: |
| raise ValueError( |
| f"GroupLayerNorm expected trailing shape ({self.n_groups}, {self.d_group}), got {x.shape[-2:]}." |
| ) |
| if pathway is None: |
| pathway = "even" |
| from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype |
| from hamiltonzero.model.tree import _kfac_name_kw |
|
|
| out_cdtype = ( |
| _compute_dtype() if pathway in ("even", "hypernet_eside") else jnp.float32 |
| ) |
| stats_dtype = jnp.promote_types(jnp.float32, x.dtype) |
| x_hi = x.astype(stats_dtype) if x.dtype != stats_dtype else x |
| ms = jnp.mean(x_hi * x_hi, axis=-1, keepdims=True) |
| normalized_hi = x_hi * jax.lax.rsqrt(ms + self.eps) |
| normalized = ( |
| normalized_hi.astype(out_cdtype) |
| if normalized_hi.dtype != out_cdtype |
| else normalized_hi |
| ) |
| w = ( |
| self.weight.astype(out_cdtype) |
| if self.weight.dtype != out_cdtype |
| else self.weight |
| ) |
| y = normalized * w |
| from hamiltonzero.optim.blocks import ( |
| register_structural_trailing_stacked_scale_and_shift, |
| ) |
|
|
| return register_structural_trailing_stacked_scale_and_shift( |
| y, |
| normalized, |
| kfac_structural_mask, |
| self.weight, |
| repeat_ndim=kfac_repeat_ndim, |
| context_primal_reused_over_walkers=kfac_context_primal_reused_over_walkers, |
| **_kfac_name_kw(self._use_id), |
| ) |
|
|
|
|
| class SystemFeaturizer(eqx.Module): |
| bond_group_w1: Linear |
| bond_group_w2: Linear |
| bond_group_post: Linear |
| bond_group_out: Linear |
| bond_group_ln: GroupLayerNorm |
| Q_col: Float[Array, "n_heads head_dim"] |
| K_col: BiasFreeLinear |
| V_col: BiasFreeLinear |
| ln_bond: _RMS |
| Q_row: Float[Array, "n_global_q n_heads head_dim"] |
| K_row: BiasFreeLinear |
| V_row: BiasFreeLinear |
| ln_local_pre_row: _RMS |
| ln_edge_cond: _RMS |
| edge_cond_w1: Linear |
| edge_cond_w2: Linear |
| ln_edge_global: _RMS |
| edge_global_film: Linear |
| edge_residual_proj: BiasFreeLinear |
| ln_edge_out: _RMS |
| ln_local_out: _RMS |
| ln_global_out: _RMS |
| zeeman_w1: Linear |
| zeeman_group_proj: Linear |
| zeeman_group_ln: GroupLayerNorm |
| zeeman_group_post: Linear |
| zeeman_w2: Linear |
| Q_row2: Float[Array, "n_global_q n_heads head_dim"] |
| K_row2: BiasFreeLinear |
| V_row2: BiasFreeLinear |
| ln_local_pre_row2: _RMS |
| ln_g_edge: _RMS |
| ln_g_zee: _RMS |
| global_w1: Linear |
| global_w2: Linear |
| ln_c1: _RMS |
| ln_c2: _RMS |
| ln_c3: _RMS |
| combine_w1: Linear |
| combine_w2: Linear |
| tok_bond_gln: Float[Array, "gb db"] |
| tok_zeeman_gln: Float[Array, "gz dz"] |
| tok_bond_key: Float[Array, "d_bond"] |
| tok_field_row: Float[Array, "d_local"] |
| tok_field_global: Float[Array, "d_global"] |
| tok_field_combine: Float[Array, "d_local"] |
| _use_id_Q_col: str = eqx.field(static=True, default="") |
| _use_id_Q_row: str = eqx.field(static=True, default="") |
| _use_id_Q_row2: str = eqx.field(static=True, default="") |
| _use_id_tok_bond_gln: str = eqx.field(static=True, default="") |
| _use_id_tok_zeeman_gln: str = eqx.field(static=True, default="") |
| _use_id_tok_bond_key: str = eqx.field(static=True, default="") |
| _use_id_tok_field_row: str = eqx.field(static=True, default="") |
| _use_id_tok_field_global: str = eqx.field(static=True, default="") |
| _use_id_tok_field_combine: str = eqx.field(static=True, default="") |
| d_bond: int = eqx.field(static=True) |
| d_local: int = eqx.field(static=True) |
| d_global: int = eqx.field(static=True) |
| d_edge: int = eqx.field(static=True) |
| n_heads: int = eqx.field(static=True) |
| head_dim: int = eqx.field(static=True) |
| n_global_q: int = eqx.field(static=True) |
| d_hidden_edge: int = eqx.field(static=True) |
| polar_group_norm_tau: float = eqx.field(static=True, default=0.001) |
| polar_group_norm_bond_hidden: int = eqx.field(static=True, default=128) |
| polar_group_norm_n_bond_groups: int = eqx.field(static=True, default=16) |
| polar_group_norm_d_bond_group: int = eqx.field(static=True, default=16) |
| polar_group_norm_n_zeeman_groups: int = eqx.field(static=True, default=16) |
| polar_group_norm_d_zeeman_group: int = eqx.field(static=True, default=16) |
|
|
| def __init__( |
| self, |
| *, |
| key: PRNGKeyArray, |
| d_bond: int, |
| n_heads: int, |
| head_dim: int, |
| n_global_q: int, |
| d_edge: int, |
| d_hidden_edge: int, |
| polar_group_norm_tau: float, |
| polar_group_norm_bond_hidden: int, |
| polar_group_norm_n_bond_groups: int, |
| polar_group_norm_d_bond_group: int, |
| polar_group_norm_n_zeeman_groups: int, |
| polar_group_norm_d_zeeman_group: int, |
| zeeman_hidden_dim: int, |
| global_hidden_dim: int, |
| combine_hidden_dim: int, |
| token_initial_scale: float, |
| ): |
| d_local = n_heads * head_dim |
| d_global = n_global_q * n_heads * head_dim |
| n_h_kernel = 2 * n_heads |
| d_local_kernel = n_h_kernel * head_dim |
| edge_cond_in = d_bond + 2 * d_local |
| self.polar_group_norm_tau = float(polar_group_norm_tau) |
| self.polar_group_norm_bond_hidden = int(polar_group_norm_bond_hidden) |
| self.polar_group_norm_n_bond_groups = int(polar_group_norm_n_bond_groups) |
| self.polar_group_norm_d_bond_group = int(polar_group_norm_d_bond_group) |
| self.polar_group_norm_n_zeeman_groups = int(polar_group_norm_n_zeeman_groups) |
| self.polar_group_norm_d_zeeman_group = int(polar_group_norm_d_zeeman_group) |
| keys = jax.random.split(key, 25) |
| bond_group_dim = polar_group_norm_n_bond_groups * polar_group_norm_d_bond_group |
| self.bond_group_w1 = Linear(20, polar_group_norm_bond_hidden, key=keys[0]) |
| self.bond_group_w2 = Linear( |
| polar_group_norm_bond_hidden, bond_group_dim, key=keys[1] |
| ) |
| self.bond_group_post = Linear( |
| bond_group_dim, polar_group_norm_bond_hidden, key=keys[2] |
| ) |
| self.bond_group_ln = GroupLayerNorm( |
| polar_group_norm_n_bond_groups, polar_group_norm_d_bond_group |
| ) |
| self.bond_group_out = Linear(polar_group_norm_bond_hidden, d_bond, key=keys[3]) |
| self.Q_col = jax.random.normal(keys[4], (n_h_kernel, head_dim)) * head_dim ** ( |
| -0.5 |
| ) |
| self.K_col = BiasFreeLinear(d_bond, d_local_kernel, key=keys[5]) |
| self.V_col = BiasFreeLinear(d_bond, d_local_kernel, key=keys[6]) |
| self.ln_bond = _RMS(d_bond) |
| self.Q_row = jax.random.normal( |
| keys[7], (n_global_q, n_h_kernel, head_dim) |
| ) * head_dim ** (-0.5) |
| self.K_row = BiasFreeLinear(d_local, d_local_kernel, key=keys[8]) |
| self.V_row = BiasFreeLinear(d_local, d_local_kernel, key=keys[9]) |
| self.ln_local_pre_row = _RMS(d_local) |
| self.ln_edge_cond = _RMS(edge_cond_in) |
| self.edge_cond_w1 = Linear(edge_cond_in, d_hidden_edge, key=keys[10]) |
| self.edge_cond_w2 = Linear(d_hidden_edge, d_edge, key=keys[11]) |
| self.ln_edge_global = _RMS(d_global) |
| self.edge_global_film = Linear(d_global, 2 * d_edge, key=keys[12]) |
| self.edge_residual_proj = BiasFreeLinear(d_bond, d_edge, key=keys[13]) |
| zeeman_group_dim = ( |
| polar_group_norm_n_zeeman_groups * polar_group_norm_d_zeeman_group |
| ) |
| self.zeeman_w1 = Linear(7, zeeman_hidden_dim, key=keys[14]) |
| self.zeeman_group_proj = Linear( |
| zeeman_hidden_dim, zeeman_group_dim, key=keys[15] |
| ) |
| self.zeeman_group_ln = GroupLayerNorm( |
| polar_group_norm_n_zeeman_groups, polar_group_norm_d_zeeman_group |
| ) |
| self.zeeman_group_post = Linear( |
| zeeman_group_dim, zeeman_hidden_dim, key=keys[16] |
| ) |
| self.zeeman_w2 = Linear(zeeman_hidden_dim, d_local, key=keys[17]) |
| self.Q_row2 = jax.random.normal( |
| keys[18], (n_global_q, n_h_kernel, head_dim) |
| ) * head_dim ** (-0.5) |
| self.K_row2 = BiasFreeLinear(d_local, d_local_kernel, key=keys[19]) |
| self.V_row2 = BiasFreeLinear(d_local, d_local_kernel, key=keys[20]) |
| self.ln_local_pre_row2 = _RMS(d_local) |
| self.ln_g_edge = _RMS(d_global) |
| self.ln_g_zee = _RMS(d_global) |
| self.global_w1 = Linear(2 * d_global + 8, global_hidden_dim, key=keys[21]) |
| self.global_w2 = Linear(global_hidden_dim, d_global, key=keys[22]) |
| combine_in = 2 * d_local + d_global |
| self.ln_c1 = _RMS(d_local) |
| self.ln_c2 = _RMS(d_local) |
| self.ln_c3 = _RMS(d_global) |
| self.combine_w1 = Linear(combine_in, combine_hidden_dim, key=keys[23]) |
| self.combine_w2 = Linear(combine_hidden_dim, d_local, key=keys[24]) |
| self.ln_edge_out = _RMS(d_edge) |
| self.ln_local_out = _RMS(d_local) |
| self.ln_global_out = _RMS(d_global) |
| tok_keys = jax.random.split(jax.random.fold_in(key, 7389448), 6) |
| self.tok_bond_gln = token_initial_scale * jax.random.normal( |
| tok_keys[0], (polar_group_norm_n_bond_groups, polar_group_norm_d_bond_group) |
| ) |
| self.tok_zeeman_gln = token_initial_scale * jax.random.normal( |
| tok_keys[1], |
| (polar_group_norm_n_zeeman_groups, polar_group_norm_d_zeeman_group), |
| ) |
| self.tok_bond_key = token_initial_scale * jax.random.normal( |
| tok_keys[2], (d_bond,) |
| ) |
| self.tok_field_row = token_initial_scale * jax.random.normal( |
| tok_keys[3], (d_local,) |
| ) |
| self.tok_field_global = token_initial_scale * jax.random.normal( |
| tok_keys[4], (d_global,) |
| ) |
| self.tok_field_combine = token_initial_scale * jax.random.normal( |
| tok_keys[5], (d_local,) |
| ) |
| self.d_bond = d_bond |
| self.d_local = d_local |
| self.d_global = d_global |
| self.d_edge = d_edge |
| self.n_heads = n_heads |
| self.head_dim = head_dim |
| self.n_global_q = n_global_q |
| self.d_hidden_edge = d_hidden_edge |
|
|
| def _polar_split( |
| self, x: Float[Array, "... d"] |
| ) -> tuple[Float[Array, "... 1"], Float[Array, "... d"]]: |
| x_f32 = x.astype(jnp.float32) |
| sq = jnp.sum(x_f32 * x_f32, axis=-1, keepdims=True) |
| tau = jnp.asarray(self.polar_group_norm_tau, dtype=jnp.float32) |
| r = jnp.sqrt(sq).astype(x.dtype) |
| direction = (x_f32 * jax.lax.rsqrt(sq + tau * tau)).astype(x.dtype) |
| return (r, direction) |
|
|
| def _bond_input( |
| self, J_double_prime: Float[Array, "n n 10"] |
| ) -> Float[Array, "n n 20"]: |
| J9 = J_double_prime[..., :9] |
| eye = J_double_prime[..., 9:] |
| rJ, uJ = self._polar_split(J9) |
| return jnp.concatenate([J9, rJ, uJ, eye], axis=-1) |
|
|
| def _zeeman_input(self, h_prime: Float[Array, "n 3"]) -> Float[Array, "n d_in"]: |
| rh, uh = self._polar_split(h_prime) |
| return jnp.concatenate([h_prime, rh, uh], axis=-1) |
|
|
| def _embed_bonds( |
| self, |
| J_double_prime: Float[Array, "n n 10"], |
| *, |
| pathway: str, |
| structural_mask: Float[Array, "n n"] | None = None, |
| ) -> Float[Array, "n n d_bond"]: |
| if structural_mask is None: |
| structural_mask = jnp.ones(J_double_prime.shape[:-1], dtype=bool) |
| z = fused_silu( |
| self._dense_structural( |
| self.bond_group_w1, |
| self._bond_input(J_double_prime), |
| structural_mask, |
| repeat_ndim=2, |
| pathway=pathway, |
| ) |
| ) |
| z = self._dense_structural( |
| self.bond_group_w2, z, structural_mask, repeat_ndim=2, pathway=pathway |
| ) |
| pair_shape = J_double_prime.shape[:-1] |
| z = z.reshape( |
| *pair_shape, |
| self.polar_group_norm_n_bond_groups, |
| self.polar_group_norm_d_bond_group, |
| ) |
| z = self.bond_group_ln( |
| z, |
| pathway=pathway, |
| kfac_structural_mask=structural_mask, |
| kfac_repeat_ndim=2, |
| kfac_context_primal_reused_over_walkers=True, |
| ) |
| from hamiltonzero.optim.spin_blocks import register_small_full |
|
|
| present = jnp.any(J_double_prime[..., :9] != 0, axis=-1) | ( |
| J_double_prime[..., 9] > 0.5 |
| ) |
| tok = register_small_full( |
| self.tok_bond_gln, tag_id=self._use_id_tok_bond_gln |
| ).astype(z.dtype) |
| z = jnp.where(present[..., None, None], z, tok) |
| z = z.reshape(*pair_shape, -1) |
| z = fused_silu( |
| self._dense_structural( |
| self.bond_group_post, z, structural_mask, repeat_ndim=2, pathway=pathway |
| ) |
| ) |
| return fused_silu( |
| self._dense_structural( |
| self.bond_group_out, z, structural_mask, repeat_ndim=2, pathway=pathway |
| ) |
| ) |
|
|
| def _embed_zeeman( |
| self, |
| h_prime: Float[Array, "n 3"], |
| *, |
| pathway: str, |
| structural_mask: Float[Array, "n"] | None = None, |
| ) -> Float[Array, "n d_local"]: |
| if structural_mask is None: |
| structural_mask = jnp.ones(h_prime.shape[:-1], dtype=bool) |
| z = fused_silu( |
| self._dense_structural( |
| self.zeeman_w1, |
| self._zeeman_input(h_prime), |
| structural_mask, |
| repeat_ndim=1, |
| pathway=pathway, |
| ) |
| ) |
| z = self._dense_structural( |
| self.zeeman_group_proj, z, structural_mask, repeat_ndim=1, pathway=pathway |
| ) |
| n = h_prime.shape[0] |
| z = z.reshape( |
| n, |
| self.polar_group_norm_n_zeeman_groups, |
| self.polar_group_norm_d_zeeman_group, |
| ) |
| z = self.zeeman_group_ln( |
| z, |
| pathway=pathway, |
| kfac_structural_mask=structural_mask, |
| kfac_repeat_ndim=1, |
| kfac_context_primal_reused_over_walkers=True, |
| ) |
| from hamiltonzero.optim.spin_blocks import register_small_full |
|
|
| present = jnp.any(h_prime != 0, axis=-1) |
| tok = register_small_full( |
| self.tok_zeeman_gln, tag_id=self._use_id_tok_zeeman_gln |
| ).astype(z.dtype) |
| z = jnp.where(present[:, None, None], z, tok) |
| z = z.reshape(n, -1) |
| z = fused_silu( |
| self._dense_structural( |
| self.zeeman_group_post, |
| z, |
| structural_mask, |
| repeat_ndim=1, |
| pathway=pathway, |
| ) |
| ) |
| return self._dense_structural( |
| self.zeeman_w2, z, structural_mask, repeat_ndim=1, pathway=pathway |
| ) |
|
|
| def _bond_norm_gated( |
| self, |
| bond_emb: Float[Array, "n n d_bond"], |
| J_double_prime: Float[Array, "n n 10"], |
| *, |
| pathway: str, |
| structural_mask: Float[Array, "n n"] | None = None, |
| ) -> Float[Array, "n n d_bond"]: |
| from hamiltonzero.optim.spin_blocks import register_small_full |
|
|
| present = jnp.any(J_double_prime[..., :9] != 0, axis=-1) | ( |
| J_double_prime[..., 9] > 0.5 |
| ) |
| if structural_mask is None: |
| structural_mask = jnp.ones(present.shape, dtype=bool) |
| tok = register_small_full( |
| self.tok_bond_key, tag_id=self._use_id_tok_bond_key |
| ).astype(bond_emb.dtype) |
| return jnp.where( |
| present[..., None], |
| self._norm_structural( |
| self.ln_bond, bond_emb, structural_mask, repeat_ndim=2, pathway=pathway |
| ), |
| tok, |
| ) |
|
|
| @staticmethod |
| def _dense_structural(lin, x, structural_mask, *, repeat_ndim: int, pathway: str): |
| return lin( |
| x, |
| pathway=pathway, |
| kfac_structural_mask=structural_mask, |
| kfac_repeat_ndim=repeat_ndim, |
| kfac_context_primal_reused_over_walkers=True, |
| ) |
|
|
| @staticmethod |
| def _norm_structural(norm, x, structural_mask, *, repeat_ndim: int, pathway: str): |
| return norm( |
| x, |
| pathway=pathway, |
| kfac_structural_mask=structural_mask, |
| kfac_repeat_ndim=repeat_ndim, |
| kfac_context_primal_reused_over_walkers=True, |
| ) |
|
|
| def _dense_bare( |
| self, lin: Linear, x: Float[Array, "..."], *, pathway: str |
| ) -> Float[Array, "..."]: |
| from hamiltonzero.model.tree import _tagged_dense |
|
|
| return _tagged_dense( |
| lin.weight, |
| lin.bias, |
| x, |
| tag_id=getattr(lin, "_use_id", ""), |
| pathway=pathway, |
| kfac_structural_mask=jnp.asarray(True), |
| kfac_scan_shared=False, |
| kfac_repeat_ndim=0, |
| kfac_context_primal_reused_over_walkers=True, |
| ) |
|
|
| @staticmethod |
| def _eval_tiles(n: int, tile_size: int): |
| tile_size = int(tile_size) |
| if tile_size < 1: |
| raise ValueError(f"tile_size must be positive, got {tile_size}") |
| return tuple((slice(j, min(j + tile_size, n)) for j in range(0, n, tile_size))) |
|
|
| def eval_embed_local_rows( |
| self, |
| J_double_prime_rows: Float[Array, "r n 10"], |
| row_mask: Float[Array, "r"], |
| mask: Float[Array, "n"], |
| *, |
| tile_size: int = 128, |
| ) -> tuple[Float[Array, "r n d_bond"], Float[Array, "r d_local"]]: |
| if J_double_prime_rows.ndim != 3 or J_double_prime_rows.shape[-1] != 10: |
| raise ValueError( |
| f"J_double_prime_rows must have shape [R,N,10], got {J_double_prime_rows.shape}" |
| ) |
| r, n = J_double_prime_rows.shape[:2] |
| if row_mask.shape != (r,) or mask.shape != (n,): |
| raise ValueError( |
| f"row_mask/mask must match J rows/columns, got {row_mask.shape}, {mask.shape}, {J_double_prime_rows.shape}" |
| ) |
| from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype |
| from hamiltonzero.optim.spin_blocks import register_small_full |
|
|
| PW = "even" |
| dtype = _compute_dtype() |
| mr = row_mask.astype(dtype) |
| mc = mask.astype(dtype) |
| pair_mask = mr[:, None] * mc[None, :] |
| bond_tiles = [] |
| tiles = self._eval_tiles(n, tile_size) |
| for sl in tiles: |
| pm = pair_mask[:, sl] |
| bond_tile = self._embed_bonds( |
| J_double_prime_rows[:, sl], pathway=PW, structural_mask=pm |
| ) |
| bond_tiles.append(bond_tile * pm[..., None]) |
| bond_emb = jnp.concatenate(bond_tiles, axis=1) |
| n_h_kernel = 2 * self.n_heads |
| scale = self.head_dim ** (-0.5) |
| Q_col = register_small_full(self.Q_col, tag_id=self._use_id_Q_col).astype(dtype) |
| score_max = jnp.full((r, n_h_kernel), -jnp.inf, dtype=dtype) |
| for sl in tiles: |
| pm = pair_mask[:, sl] |
| bn = self._bond_norm_gated( |
| bond_emb[:, sl], |
| J_double_prime_rows[:, sl], |
| pathway=PW, |
| structural_mask=pm, |
| ) |
| k_tile = self._dense_structural( |
| self.K_col, bn, pm, repeat_ndim=2, pathway=PW |
| ).reshape(r, sl.stop - sl.start, n_h_kernel, self.head_dim) |
| scores = jnp.einsum("hd,ijhd->ijh", Q_col, k_tile) * scale |
| scores = jnp.where( |
| pm[..., None] > 0, |
| scores, |
| jnp.asarray(-1000000000.0, dtype=scores.dtype), |
| ) |
| score_max = jnp.maximum(score_max, jnp.max(scores, axis=1)) |
| denom = jnp.zeros((r, n_h_kernel), dtype=dtype) |
| numer = jnp.zeros((r, n_h_kernel, self.head_dim), dtype=dtype) |
| for sl in tiles: |
| width = sl.stop - sl.start |
| pm = pair_mask[:, sl] |
| bn = self._bond_norm_gated( |
| bond_emb[:, sl], |
| J_double_prime_rows[:, sl], |
| pathway=PW, |
| structural_mask=pm, |
| ) |
| k_tile = self._dense_structural( |
| self.K_col, bn, pm, repeat_ndim=2, pathway=PW |
| ).reshape(r, width, n_h_kernel, self.head_dim) |
| scores = jnp.einsum("hd,ijhd->ijh", Q_col, k_tile) * scale |
| scores = jnp.where( |
| pm[..., None] > 0, |
| scores, |
| jnp.asarray(-1000000000.0, dtype=scores.dtype), |
| ) |
| weight = jnp.exp(scores - score_max[:, None, :]) |
| v_tile = self._dense_structural( |
| self.V_col, bn, pm, repeat_ndim=2, pathway=PW |
| ).reshape(r, width, n_h_kernel, self.head_dim) |
| denom = denom + jnp.sum(weight, axis=1) |
| numer = numer + jnp.einsum("ijh,ijhd->ihd", weight, v_tile) |
| col_out = numer / jnp.maximum( |
| denom[..., None], jnp.asarray(1e-30, dtype=numer.dtype) |
| ) |
| gate = col_out[:, : self.n_heads, :] |
| val = col_out[:, self.n_heads :, :] |
| local_desc_rows = (jax.nn.sigmoid(gate) * val).reshape(r, self.d_local) * mr[ |
| :, None |
| ] |
| return (bond_emb, local_desc_rows) |
|
|
| def eval_jh_stats_rows( |
| self, |
| J_double_prime_rows: Float[Array, "r n 10"], |
| row_mask: Float[Array, "r"], |
| mask: Float[Array, "n"], |
| *, |
| row_indices: Int[Array, "r"] | None = None, |
| ) -> tuple[Float[Array, ""], Float[Array, ""]]: |
| r, n = J_double_prime_rows.shape[:2] |
| if row_indices is None: |
| if r != n: |
| raise ValueError("row_indices is required when R != N") |
| row_indices = jnp.arange(n, dtype=jnp.int32) |
| row_indices = jnp.asarray(row_indices, dtype=jnp.int32) |
| col_indices = jnp.arange(n, dtype=jnp.int32) |
| active = ( |
| row_mask[:, None].astype(J_double_prime_rows.dtype) |
| * mask[None, :].astype(J_double_prime_rows.dtype) |
| * (row_indices[:, None] != col_indices[None, :]).astype( |
| J_double_prime_rows.dtype |
| ) |
| ) |
| norm2 = jnp.sum(jnp.square(J_double_prime_rows[..., :9]), axis=-1) |
| return (jnp.sum(norm2 * active), jnp.sum(active)) |
|
|
| def eval_edge_rows( |
| self, |
| *, |
| bond_emb_rows: Float[Array, "r n d_bond"], |
| local_rows: Float[Array, "r d_local"], |
| local_final_all: Float[Array, "n d_local"], |
| global_feat: Float[Array, "d_global"], |
| row_indices: Int[Array, "r"], |
| mask: Float[Array, "n"], |
| tile_size: int = 128, |
| ) -> Float[Array, "r n d_edge"]: |
| PW = "even" |
| r, n = bond_emb_rows.shape[:2] |
| row_indices = jnp.asarray(row_indices, dtype=jnp.int32) |
| if row_indices.shape != (r,): |
| raise ValueError("row_indices must have shape [R]") |
| if local_rows.shape != (r, self.d_local): |
| raise ValueError("local_rows must have shape [R,d_local]") |
| if local_final_all.shape != (n, self.d_local): |
| raise ValueError("local_final_all must have shape [N,d_local]") |
| if mask.shape != (n,): |
| raise ValueError("mask must have shape [N]") |
| from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype |
|
|
| dtype = _compute_dtype() |
| m = mask.astype(dtype) |
| mr = m[row_indices] |
| pair_mask = mr[:, None] * m[None, :] |
| global_mask = jnp.asarray(True) |
| film_in = self._norm_structural( |
| self.ln_edge_global, global_feat, global_mask, repeat_ndim=0, pathway=PW |
| ) |
| film = self._dense_bare(self.edge_global_film, film_in, pathway=PW) |
| gamma, beta = jnp.split(film, 2, axis=-1) |
| gamma = 0.1 * jnp.tanh(gamma) |
| beta = 0.1 * beta |
| edge_tiles = [] |
| for sl in self._eval_tiles(n, tile_size): |
| pm = pair_mask[:, sl] |
| bond = bond_emb_rows[:, sl] |
| width = sl.stop - sl.start |
| li = jnp.broadcast_to(local_rows[:, None, :], (r, width, self.d_local)) |
| lj = jnp.broadcast_to( |
| local_final_all[sl][None, :, :], (r, width, self.d_local) |
| ) |
| edge_in = jnp.concatenate([bond, li, lj], axis=-1) |
| edge_in = self._norm_structural( |
| self.ln_edge_cond, edge_in, pm, repeat_ndim=2, pathway=PW |
| ) |
| core = self._dense_structural( |
| self.edge_cond_w2, |
| fused_silu( |
| self._dense_structural( |
| self.edge_cond_w1, edge_in, pm, repeat_ndim=2, pathway=PW |
| ) |
| ), |
| pm, |
| repeat_ndim=2, |
| pathway=PW, |
| ) |
| update_edge = core * (1.0 + gamma[None, None, :]) + beta[None, None, :] |
| residual = self._dense_structural( |
| self.edge_residual_proj, bond, pm, repeat_ndim=2, pathway=PW |
| ) |
| edge_tile = ( |
| self._norm_structural( |
| self.ln_edge_out, |
| residual + update_edge, |
| pm, |
| repeat_ndim=2, |
| pathway=PW, |
| ) |
| * pm[..., None] |
| ) |
| edge_tiles.append(edge_tile) |
| return jnp.concatenate(edge_tiles, axis=1) |
|
|
| def eval_finalize_local_rows( |
| self, |
| *, |
| J_double_prime_rows: Float[Array, "r n 10"], |
| local_desc_rows: Float[Array, "r d_local"], |
| local_desc_all: Float[Array, "n d_local"], |
| row_indices: Int[Array, "r"], |
| mask: Float[Array, "n"], |
| h_prime: Float[Array, "n 3"], |
| jh_stats: tuple[Float[Array, ""], Float[Array, ""]] | None = None, |
| ) -> tuple[Float[Array, "r d_local"], Float[Array, "d_global"]]: |
| from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype |
| from hamiltonzero.optim.spin_blocks import register_small_full |
|
|
| PW = "even" |
| dtype = _compute_dtype() |
| r, n = J_double_prime_rows.shape[:2] |
| row_indices = jnp.asarray(row_indices, dtype=jnp.int32) |
| if row_indices.shape != (r,): |
| raise ValueError("row_indices must have shape [R]") |
| if local_desc_rows.shape != (r, self.d_local): |
| raise ValueError("local_desc_rows has the wrong shape") |
| if local_desc_all.shape != (n, self.d_local) or mask.shape != (n,): |
| raise ValueError("local_desc_all/mask must have global width N") |
| m = mask.astype(dtype) |
| mr = m[row_indices] |
| n_h_kernel = 2 * self.n_heads |
| scale = self.head_dim ** (-0.5) |
| global_mask = jnp.asarray(True) |
|
|
| def row_descriptor(x, ln_pre, Q_arr, K_lin, V_lin, q_use_id, present=None): |
| x_norm = self._norm_structural(ln_pre, x, m, repeat_ndim=1, pathway=PW) |
| if present is not None: |
| tok_row = register_small_full( |
| self.tok_field_row, tag_id=self._use_id_tok_field_row |
| ).astype(x_norm.dtype) |
| x_norm = jnp.where(present[:, None], x_norm, tok_row) |
| qr = register_small_full(Q_arr, tag_id=q_use_id).astype(dtype) |
| kr = self._dense_structural( |
| K_lin, x_norm, m, repeat_ndim=1, pathway=PW |
| ).reshape(n, n_h_kernel, self.head_dim) |
| vr = self._dense_structural( |
| V_lin, x_norm, m, repeat_ndim=1, pathway=PW |
| ).reshape(n, n_h_kernel, self.head_dim) |
| sc = jnp.einsum("qhd,ihd->qhi", qr, kr) * scale |
| sc = jnp.where( |
| m[None, None, :] > 0, sc, jnp.asarray(-1000000000.0, dtype=sc.dtype) |
| ) |
| a = jax.nn.softmax(sc, axis=-1) |
| out = jnp.einsum("qhi,ihd->qhd", a, vr) |
| out = jax.nn.sigmoid(out[:, : self.n_heads, :]) * out[:, self.n_heads :, :] |
| return out.reshape(self.d_global) |
|
|
| h_prime = h_prime.astype(dtype) |
| local_prime_all = ( |
| self._embed_zeeman(h_prime, pathway=PW, structural_mask=m) * m[:, None] |
| ) |
| local_prime_rows = local_prime_all[row_indices] |
| h_present = jnp.any(h_prime != 0, axis=-1) |
| g_edge = row_descriptor( |
| local_desc_all, |
| self.ln_local_pre_row, |
| self.Q_row, |
| self.K_row, |
| self.V_row, |
| self._use_id_Q_row, |
| ) |
| g_zee = row_descriptor( |
| local_prime_all, |
| self.ln_local_pre_row2, |
| self.Q_row2, |
| self.K_row2, |
| self.V_row2, |
| self._use_id_Q_row2, |
| present=h_present, |
| ) |
| if jh_stats is None: |
| if r != n: |
| raise ValueError( |
| "row-sharded featurization requires the psum result from eval_jh_stats_rows" |
| ) |
| jh_stats = self.eval_jh_stats_rows( |
| J_double_prime_rows, mr, m, row_indices=row_indices |
| ) |
| sum_j2, count_j = jh_stats |
| s_j2 = sum_j2 / jnp.maximum(count_j, 1.0) |
| s_h2 = jnp.sum(jnp.sum(jnp.square(h_prime), axis=-1) * m) / jnp.maximum( |
| jnp.sum(m), 1.0 |
| ) |
|
|
| def safe_sqrt(x): |
| ok = x > 0 |
| return jnp.where(ok, jnp.sqrt(jnp.where(ok, x, 1.0)), 0.0) |
|
|
| rj = safe_sqrt(s_j2) |
| rh = safe_sqrt(s_h2) |
| ok = rj + rh > 0 |
| theta = jnp.arctan2(jnp.where(ok, rh, 0.0), jnp.where(ok, rj, 1.0)) |
| jh_features = jnp.stack( |
| [ |
| jnp.log1p(rj), |
| jnp.log1p(rh), |
| jnp.sin(2.0 * theta), |
| jnp.cos(2.0 * theta), |
| jnp.sin(4.0 * theta), |
| jnp.cos(4.0 * theta), |
| jnp.sin(8.0 * theta), |
| jnp.cos(8.0 * theta), |
| ] |
| ).astype(dtype) |
| global_cat = jnp.concatenate( |
| [ |
| self._norm_structural( |
| self.ln_g_edge, g_edge, global_mask, repeat_ndim=0, pathway=PW |
| ), |
| jnp.where( |
| jnp.any(h_present), |
| self._norm_structural( |
| self.ln_g_zee, g_zee, global_mask, repeat_ndim=0, pathway=PW |
| ), |
| register_small_full( |
| self.tok_field_global, tag_id=self._use_id_tok_field_global |
| ).astype(dtype), |
| ), |
| jh_features, |
| ], |
| axis=-1, |
| ) |
| global_raw = self._dense_bare( |
| self.global_w2, |
| fused_silu(self._dense_bare(self.global_w1, global_cat, pathway=PW)), |
| pathway=PW, |
| ) |
| global_feat = self._norm_structural( |
| self.ln_global_out, global_raw, global_mask, repeat_ndim=0, pathway=PW |
| ) |
| combine = jnp.concatenate( |
| [ |
| jnp.where( |
| h_present[row_indices, None], |
| self._norm_structural( |
| self.ln_c1, local_prime_rows, mr, repeat_ndim=1, pathway=PW |
| ), |
| register_small_full( |
| self.tok_field_combine, tag_id=self._use_id_tok_field_combine |
| ).astype(dtype), |
| ), |
| self._norm_structural( |
| self.ln_c2, local_desc_rows, mr, repeat_ndim=1, pathway=PW |
| ), |
| jnp.broadcast_to( |
| self._norm_structural( |
| self.ln_c3, global_feat, global_mask, repeat_ndim=0, pathway=PW |
| ), |
| (r, self.d_global), |
| ), |
| ], |
| axis=-1, |
| ) |
| update = self._dense_structural( |
| self.combine_w2, |
| fused_silu( |
| self._dense_structural( |
| self.combine_w1, combine, mr, repeat_ndim=1, pathway=PW |
| ) |
| ), |
| mr, |
| repeat_ndim=1, |
| pathway=PW, |
| ) |
| local_rows = ( |
| self._norm_structural( |
| self.ln_local_out, |
| local_prime_rows + update, |
| mr, |
| repeat_ndim=1, |
| pathway=PW, |
| ) |
| * mr[:, None] |
| ) |
| return (local_rows, global_feat) |
|
|
| def eval_streamed( |
| self, |
| J_double_prime: Float[Array, "n n 10"], |
| mask: Float[Array, "n"], |
| h_prime: Float[Array, "n 3"], |
| *, |
| tile_size: int = 128, |
| ): |
| n = J_double_prime.shape[0] |
| idx = jnp.arange(n, dtype=jnp.int32) |
| bond, local_desc = self.eval_embed_local_rows( |
| J_double_prime, mask, mask, tile_size=tile_size |
| ) |
| stats = self.eval_jh_stats_rows(J_double_prime, mask, mask, row_indices=idx) |
| local, global_feat = self.eval_finalize_local_rows( |
| J_double_prime_rows=J_double_prime, |
| local_desc_rows=local_desc, |
| local_desc_all=local_desc, |
| row_indices=idx, |
| mask=mask, |
| h_prime=h_prime, |
| jh_stats=stats, |
| ) |
| edge = self.eval_edge_rows( |
| bond_emb_rows=bond, |
| local_rows=local, |
| local_final_all=local, |
| global_feat=global_feat, |
| row_indices=idx, |
| mask=mask, |
| tile_size=tile_size, |
| ) |
| return (edge, local, global_feat) |
|
|
| def __call__( |
| self, |
| J_double_prime: Float[Array, "n n 10"], |
| mask: Float[Array, "n"], |
| h_prime: Float[Array, "n 3"], |
| ) -> tuple[ |
| Float[Array, "n n d_edge"], Float[Array, "n d_local"], Float[Array, "d_global"] |
| ]: |
| PW = "even" |
| n = J_double_prime.shape[0] |
| from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype |
|
|
| dtype = _compute_dtype() |
| m = mask.astype(dtype) |
| pair_mask = m[:, None] * m[None, :] |
| global_mask = jnp.asarray(True) |
| head_dim = self.head_dim |
| n_heads = self.n_heads |
| n_global_q = self.n_global_q |
| scale = head_dim ** (-0.5) |
| bond_emb = self._embed_bonds( |
| J_double_prime, pathway=PW, structural_mask=pair_mask |
| ) |
| bond_emb = bond_emb * pair_mask[..., None] |
| n_h_kernel = 2 * n_heads |
| from hamiltonzero.optim.spin_blocks import register_small_full |
|
|
| bond_norm = self._bond_norm_gated( |
| bond_emb, J_double_prime, pathway=PW, structural_mask=pair_mask |
| ) |
| Q_col = register_small_full(self.Q_col, tag_id=self._use_id_Q_col).astype(dtype) |
| K = self._dense_structural( |
| self.K_col, bond_norm, pair_mask, repeat_ndim=2, pathway=PW |
| ).reshape(n, n, n_h_kernel, head_dim) |
| V = self._dense_structural( |
| self.V_col, bond_norm, pair_mask, repeat_ndim=2, pathway=PW |
| ).reshape(n, n, n_h_kernel, head_dim) |
| scores = jnp.einsum("hd,ijhd->ijh", Q_col, K) * scale |
| scores = jnp.where( |
| pair_mask[..., None] > 0, |
| scores, |
| jnp.asarray(-1000000000.0, dtype=scores.dtype), |
| ) |
| attn = jax.nn.softmax(scores, axis=1) |
| col_out = jnp.einsum("ijh,ijhd->ihd", attn, V) |
| gate = col_out[:, :n_heads, :] |
| val = col_out[:, n_heads:, :] |
| col_out = jax.nn.sigmoid(gate) * val |
| local_i_raw = col_out.reshape(n, self.d_local) * m[:, None] |
| h_prime = h_prime.astype(dtype) |
|
|
| def _row_descriptor(x, ln_pre, Q_arr, K_lin, V_lin, q_use_id, present=None): |
| x_norm = self._norm_structural(ln_pre, x, m, repeat_ndim=1, pathway=PW) |
| if present is not None: |
| tok_row = register_small_full( |
| self.tok_field_row, tag_id=self._use_id_tok_field_row |
| ).astype(x_norm.dtype) |
| x_norm = jnp.where(present[:, None], x_norm, tok_row) |
| Qr = register_small_full(Q_arr, tag_id=q_use_id).astype(dtype) |
| Kr = self._dense_structural( |
| K_lin, x_norm, m, repeat_ndim=1, pathway=PW |
| ).reshape(n, n_h_kernel, head_dim) |
| Vr = self._dense_structural( |
| V_lin, x_norm, m, repeat_ndim=1, pathway=PW |
| ).reshape(n, n_h_kernel, head_dim) |
| sc = jnp.einsum("qhd,ihd->qhi", Qr, Kr) * scale |
| sc = jnp.where( |
| m[None, None, :] > 0, sc, jnp.asarray(-1000000000.0, dtype=sc.dtype) |
| ) |
| a = jax.nn.softmax(sc, axis=-1) |
| ro = jnp.einsum("qhi,ihd->qhd", a, Vr) |
| ro = jax.nn.sigmoid(ro[:, :n_heads, :]) * ro[:, n_heads:, :] |
| return ro.reshape(self.d_global) |
|
|
| local_desc_i = local_i_raw |
| local_i_prime = self._embed_zeeman(h_prime, pathway=PW, structural_mask=m) |
| local_i_prime = local_i_prime * m[:, None] |
| h_present = jnp.any(h_prime != 0, axis=-1) |
| g_edge = _row_descriptor( |
| local_desc_i, |
| self.ln_local_pre_row, |
| self.Q_row, |
| self.K_row, |
| self.V_row, |
| self._use_id_Q_row, |
| ) |
| g_zee = _row_descriptor( |
| local_i_prime, |
| self.ln_local_pre_row2, |
| self.Q_row2, |
| self.K_row2, |
| self.V_row2, |
| self._use_id_Q_row2, |
| present=h_present, |
| ) |
| off_diag = pair_mask * (1.0 - jnp.eye(n, dtype=dtype)) |
| bond_magnitude2 = jnp.sum( |
| jnp.square(J_double_prime[..., :9].astype(dtype)), axis=-1 |
| ) |
| mean_j2 = jnp.sum(bond_magnitude2 * off_diag) / jnp.maximum( |
| jnp.sum(off_diag), 1.0 |
| ) |
| mean_h2 = jnp.sum(jnp.sum(jnp.square(h_prime), axis=-1) * m) / jnp.maximum( |
| jnp.sum(m), 1.0 |
| ) |
|
|
| def _safe_sqrt(x): |
| ok = x > 0 |
| return jnp.where(ok, jnp.sqrt(jnp.where(ok, x, 1.0)), 0.0) |
|
|
| rj = _safe_sqrt(mean_j2) |
| rh = _safe_sqrt(mean_h2) |
| nonzero = rj + rh > 0 |
| theta = jnp.arctan2(jnp.where(nonzero, rh, 0.0), jnp.where(nonzero, rj, 1.0)) |
| jh_features = jnp.stack( |
| [ |
| jnp.log1p(rj), |
| jnp.log1p(rh), |
| jnp.sin(2.0 * theta), |
| jnp.cos(2.0 * theta), |
| jnp.sin(4.0 * theta), |
| jnp.cos(4.0 * theta), |
| jnp.sin(8.0 * theta), |
| jnp.cos(8.0 * theta), |
| ] |
| ).astype(dtype) |
| global_cat = jnp.concatenate( |
| [ |
| self._norm_structural( |
| self.ln_g_edge, g_edge, global_mask, repeat_ndim=0, pathway=PW |
| ), |
| jnp.where( |
| jnp.any(h_present), |
| self._norm_structural( |
| self.ln_g_zee, g_zee, global_mask, repeat_ndim=0, pathway=PW |
| ), |
| register_small_full( |
| self.tok_field_global, tag_id=self._use_id_tok_field_global |
| ).astype(dtype), |
| ), |
| jh_features, |
| ], |
| axis=-1, |
| ) |
| global_raw = self._dense_bare( |
| self.global_w2, |
| fused_silu(self._dense_bare(self.global_w1, global_cat, pathway=PW)), |
| pathway=PW, |
| ) |
| global_feat = self._norm_structural( |
| self.ln_global_out, global_raw, global_mask, repeat_ndim=0, pathway=PW |
| ) |
| combine_in = jnp.concatenate( |
| [ |
| jnp.where( |
| h_present[:, None], |
| self._norm_structural( |
| self.ln_c1, local_i_prime, m, repeat_ndim=1, pathway=PW |
| ), |
| register_small_full( |
| self.tok_field_combine, tag_id=self._use_id_tok_field_combine |
| ).astype(dtype), |
| ), |
| self._norm_structural( |
| self.ln_c2, local_desc_i, m, repeat_ndim=1, pathway=PW |
| ), |
| jnp.broadcast_to( |
| self._norm_structural( |
| self.ln_c3, global_feat, global_mask, repeat_ndim=0, pathway=PW |
| )[None, :], |
| (n, self.d_global), |
| ), |
| ], |
| axis=-1, |
| ) |
| local_update = self._dense_structural( |
| self.combine_w2, |
| fused_silu( |
| self._dense_structural( |
| self.combine_w1, combine_in, m, repeat_ndim=1, pathway=PW |
| ) |
| ), |
| m, |
| repeat_ndim=1, |
| pathway=PW, |
| ) |
| local_i = self._norm_structural( |
| self.ln_local_out, |
| local_i_prime + local_update, |
| m, |
| repeat_ndim=1, |
| pathway=PW, |
| ) |
| local_i = local_i * m[:, None] |
| local_i_for_i = jnp.broadcast_to(local_i[:, None, :], (n, n, self.d_local)) |
| local_i_for_j = jnp.broadcast_to(local_i[None, :, :], (n, n, self.d_local)) |
| edge_cond_in = jnp.concatenate( |
| [bond_emb, local_i_for_i, local_i_for_j], axis=-1 |
| ) |
| edge_cond_norm = self._norm_structural( |
| self.ln_edge_cond, edge_cond_in, pair_mask, repeat_ndim=2, pathway=PW |
| ) |
| edge_core = self._dense_structural( |
| self.edge_cond_w2, |
| fused_silu( |
| self._dense_structural( |
| self.edge_cond_w1, |
| edge_cond_norm, |
| pair_mask, |
| repeat_ndim=2, |
| pathway=PW, |
| ) |
| ), |
| pair_mask, |
| repeat_ndim=2, |
| pathway=PW, |
| ) |
| ln_global_for_film = self._norm_structural( |
| self.ln_edge_global, global_feat, global_mask, repeat_ndim=0, pathway=PW |
| ) |
| film = self._dense_bare(self.edge_global_film, ln_global_for_film, pathway=PW) |
| gamma, beta = jnp.split(film, 2, axis=-1) |
| gamma = 0.1 * jnp.tanh(gamma) |
| beta = 0.1 * beta |
| edge_update = edge_core * (1.0 + gamma[None, None, :]) + beta[None, None, :] |
| bond_emb_proj = self._dense_structural( |
| self.edge_residual_proj, bond_emb, pair_mask, repeat_ndim=2, pathway=PW |
| ) |
| edge_ij = bond_emb_proj + edge_update |
| edge_ij = self._norm_structural( |
| self.ln_edge_out, edge_ij, pair_mask, repeat_ndim=2, pathway=PW |
| ) |
| edge_ij = edge_ij * pair_mask[..., None] |
| return (edge_ij, local_i, global_feat) |
|
|
