|
|
import jax |
|
|
import jax.numpy as jnp |
|
|
|
|
|
from flax import linen as nn |
|
|
import jax.numpy as jnp |
|
|
|
|
|
from einops import rearrange |
|
|
|
|
|
def roll(J, shift, axis=-1): |
|
|
return jnp.roll(J, shift, axis=axis) |
|
|
|
|
|
from functools import partial |
|
|
@partial(jax.vmap, in_axes=(None, 0, None), out_axes=1) |
|
|
@partial(jax.vmap, in_axes=(None, None, 0), out_axes=1) |
|
|
def roll2d(spins, i, j): |
|
|
side = int(spins.shape[-1]**0.5) |
|
|
spins = spins.reshape(spins.shape[0], side, side) |
|
|
spins = jnp.roll(jnp.roll(spins, i, axis=-2), j, axis=-1) |
|
|
return spins.reshape(spins.shape[0], -1) |
|
|
|
|
|
class FMHA(nn.Module): |
|
|
d_model : int |
|
|
h: int |
|
|
L_eff: int |
|
|
transl_invariant: bool = True |
|
|
two_dimensional: bool = False |
|
|
|
|
|
def setup(self): |
|
|
self.v = nn.Dense(self.d_model, kernel_init=nn.initializers.xavier_uniform(), param_dtype=jnp.float64, dtype=jnp.float64) |
|
|
if self.transl_invariant: |
|
|
self.J = self.param("J", nn.initializers.xavier_uniform(), (self.h, self.L_eff), jnp.float64) |
|
|
if self.two_dimensional: |
|
|
sq_L_eff = int(self.L_eff**0.5) |
|
|
assert sq_L_eff * sq_L_eff == self.L_eff |
|
|
self.J = roll2d(self.J, jnp.arange(sq_L_eff), jnp.arange(sq_L_eff)) |
|
|
self.J = self.J.reshape(self.h, -1, self.L_eff) |
|
|
else: |
|
|
self.J = jax.vmap(roll, (None, 0), out_axes=1)(self.J, jnp.arange(self.L_eff)) |
|
|
else: |
|
|
self.J = self.param("J", nn.initializers.xavier_uniform(), (self.h, self.L_eff, self.L_eff), jnp.float64) |
|
|
|
|
|
self.W = nn.Dense(self.d_model, kernel_init=nn.initializers.xavier_uniform(), param_dtype=jnp.float64, dtype=jnp.float64) |
|
|
|
|
|
def __call__(self, x): |
|
|
v = self.v(x) |
|
|
v = rearrange(v, 'batch L_eff (h d_eff) -> batch L_eff h d_eff', h=self.h) |
|
|
v = rearrange(v, 'batch L_eff h d_eff -> batch h L_eff d_eff') |
|
|
x = jnp.matmul(self.J, v) |
|
|
x = rearrange(x, 'batch h L_eff d_eff -> batch L_eff h d_eff') |
|
|
x = rearrange(x, 'batch L_eff h d_eff -> batch L_eff (h d_eff)') |
|
|
|
|
|
x = self.W(x) |
|
|
|
|
|
return x |
|
|
|
