Upload model

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

attentions.py

@@ -0,0 +1,54 @@
+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

config.json

@@ -0,0 +1,20 @@
+{
+  "L_eff": 25,
+  "architectures": [
+    "NQSModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "vit_fnqs_config.ViTFNQSConfig",
+    "FlaxAutoModel": "vit_fnqs_model.ViTFNQSModel"
+  },
+  "b": 2,
+  "complex": true,
+  "d_model": 72,
+  "disorder": false,
+  "heads": 12,
+  "model_type": "vit_fnqs",
+  "num_layers": 4,
+  "transformers_version": "4.48.0",
+  "tras_inv": true,
+  "two_dim": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f267caaf723e8c20f1b6efbdfb4e8af3904a66d952d61cce3209e948530649a
+size 1791368

transformer_fnqs.py

@@ -0,0 +1,162 @@
+import jax
+import jax.numpy as jnp
+
+from flax import linen as nn
+import jax.numpy as jnp
+from einops import rearrange
+from .attentions import FMHA
+
+def extract_patches1d(x, b):
+    return rearrange(x, 'batch (L_eff b) -> batch L_eff b', b=b)
+
+def extract_patches2d(x, b):
+    batch = x.shape[0]
+    L_eff = int((x.shape[1] // b**2)**0.5)
+    x = x.reshape(batch, L_eff, b, L_eff, b)   # [L_eff, b, L_eff, b]
+    x = x.transpose(0, 1, 3, 2, 4)         # [L_eff, L_eff, b, b]
+    # flatten the patches
+    x = x.reshape(batch, L_eff, L_eff, -1)     # [L_eff, L_eff, b*b]
+    x = x.reshape(batch, L_eff*L_eff, -1)      # [L_eff*L_eff, b*b]
+    return x
+
+def log_cosh(x):
+    sgn_x = -2 * jnp.signbit(x.real) + 1
+    x = x * sgn_x
+    return x + jnp.log1p(jnp.exp(-2.0 * x)) - jnp.log(2.0)
+
+class Embed(nn.Module):
+    d_model : int
+    b: int
+    two_dimensional: bool = False
+
+    def setup(self):
+        if self.two_dimensional:
+            self.extract_patches = extract_patches2d
+        else:
+            self.extract_patches = extract_patches1d
+
+        self.embed = nn.Dense(self.d_model, kernel_init=nn.initializers.xavier_uniform(), param_dtype=jnp.float64, dtype=jnp.float64)
+
+    def __call__(self, x):
+        x = self.extract_patches(x, self.b)
+        x = self.embed(x)
+
+        return x
+
+class EncoderBlock(nn.Module):
+    d_model : int
+    h: int
+    L_eff: int
+    transl_invariant: bool = True
+    two_dimensional: bool = False
+
+    def setup(self):
+        self.attn = FMHA(d_model=self.d_model, h=self.h, L_eff=self.L_eff, transl_invariant=self.transl_invariant, two_dimensional=self.two_dimensional)
+        
+        self.layer_norm_1 = nn.LayerNorm(dtype=jnp.float64, param_dtype=jnp.float64)
+        self.layer_norm_2 = nn.LayerNorm(dtype=jnp.float64, param_dtype=jnp.float64)
+
+        self.ff = nn.Sequential([
+            nn.Dense(4*self.d_model, kernel_init=nn.initializers.xavier_uniform(), param_dtype=jnp.float64, dtype=jnp.float64),
+            nn.gelu,
+            nn.Dense(self.d_model, kernel_init=nn.initializers.xavier_uniform(), param_dtype=jnp.float64, dtype=jnp.float64),
+        ])
+
+    def __call__(self, x):
+        x = x + self.attn(self.layer_norm_1(x))
+
+        x = x + self.ff( self.layer_norm_2(x) )
+        return x
+    
+class Encoder(nn.Module):
+    num_layers: int
+    d_model : int
+    h: int
+    L_eff: int
+    transl_invariant: bool = True
+    two_dimensional: bool = False
+
+    def setup(self):
+        self.layers = [EncoderBlock(d_model=self.d_model, h=self.h, L_eff=self.L_eff, transl_invariant=self.transl_invariant, two_dimensional=self.two_dimensional) for _ in range(self.num_layers)]
+
+    def __call__(self, x):
+        
+        for l in self.layers:
+            x = l(x)
+
+        return x
+
+class OuputHead(nn.Module):
+    d_model : int
+    complex: bool = False
+
+    def setup(self):
+        self.out_layer_norm = nn.LayerNorm(dtype=jnp.float64, param_dtype=jnp.float64)
+
+        self.norm0 = nn.LayerNorm(use_scale=True, use_bias=True, dtype=jnp.float64, param_dtype=jnp.float64)
+        self.norm1 = nn.LayerNorm(use_scale=True, use_bias=True, dtype=jnp.float64, param_dtype=jnp.float64)
+
+        self.output_layer0 = nn.Dense(self.d_model, param_dtype=jnp.float64, dtype=jnp.float64, kernel_init=nn.initializers.xavier_uniform(), bias_init=jax.nn.initializers.zeros)
+        self.output_layer1 = nn.Dense(self.d_model, param_dtype=jnp.float64, dtype=jnp.float64, kernel_init=nn.initializers.xavier_uniform(), bias_init=jax.nn.initializers.zeros)
+
+    def __call__(self, x, return_z=False):
+
+        z = self.out_layer_norm(x.sum(axis=1))
+
+        if return_z:
+            return z
+        
+        amp = self.norm0(self.output_layer0(z))
+        
+        if self.complex:
+            sign = self.norm1(self.output_layer1(z))
+            out = amp + 1j*sign
+        else:
+            out = amp
+
+        return jnp.sum(log_cosh(out), axis=-1)
+
+class ViTFNQS(nn.Module):
+    num_layers: int
+    d_model : int
+    heads: int
+    L_eff: int
+    b: int
+    complex: bool = False
+    disorder: bool = False
+    transl_invariant: bool = True
+    two_dimensional: bool = False
+
+    def setup(self):
+        if self.disorder:
+            self.patches_and_embed = Embed(self.d_model//2, self.b, two_dimensional=self.two_dimensional)
+            self.patches_and_embed_coup = Embed(self.d_model//2, self.b, two_dimensional=self.two_dimensional)
+        else:
+            self.embed = nn.Dense(self.d_model, kernel_init=nn.initializers.xavier_uniform(), param_dtype=jnp.float64, dtype=jnp.float64)
+
+        self.encoder = Encoder(num_layers=self.num_layers, d_model=self.d_model, h=self.heads, L_eff=self.L_eff, transl_invariant=self.transl_invariant, two_dimensional=self.two_dimensional)
+
+        self.output = OuputHead(self.d_model, complex=self.complex)
+
+    def __call__(self, spins, coups, return_z=False):
+        x = jnp.atleast_2d(spins)
+
+        if self.disorder:
+            x_spins = self.patches_and_embed(x)
+            x_coups = self.patches_and_embed(coups)
+            x = jnp.concatenate((x_spins, x_coups), axis=-1)
+        else:
+            if self.two_dimensional:
+                x = extract_patches2d(x, self.b)
+            else:
+                x = extract_patches1d(x, self.b)
+            coups = jnp.broadcast_to(coups, (x.shape[0], x.shape[1], 1))
+            # coups = jnp.repeat(coups[:, None], repeats=x.shape[1], axis=1)
+            x = jnp.concatenate((x, coups), axis=-1)
+            x = self.embed(x)
+
+        x = self.encoder(x)
+
+        out = self.output(x, return_z=return_z)
+
+        return out

vit_fnqs_config.py

@@ -0,0 +1,30 @@
+from transformers import PretrainedConfig
+from typing import List
+
+
+class ViTFNQSConfig(PretrainedConfig):
+    model_type = "vit_fnqs"
+
+    def __init__(
+        self,
+        L_eff=25,
+        num_layers = 4,
+        d_model = 72,
+        heads = 12,
+        b = 2,
+        complex: bool = True,
+        disorder: bool = False,
+        tras_inv = True,
+        two_dim = True,
+        **kwargs,
+    ):
+        self.L_eff = L_eff
+        self.num_layers = num_layers
+        self.d_model = d_model
+        self.heads = heads
+        self.b = b
+        self.complex = complex
+        self.disorder = disorder
+        self.tras_inv = tras_inv
+        self.two_dim = two_dim
+        super().__init__(**kwargs)

vit_fnqs_model.py

@@ -0,0 +1,39 @@
+from transformers import FlaxPreTrainedModel
+import jax.numpy as jnp
+from .transformer_fnqs import ViTFNQS
+from .vit_fnqs_config import ViTFNQSConfig
+
+class ViTFNQSModel(FlaxPreTrainedModel):
+    config_class = ViTFNQSConfig
+
+    def __init__(
+            self, 
+            config: ViTFNQSConfig, 
+            input_shape = (jnp.zeros((1, 100)), jnp.zeros((1, 1))),
+            seed: int = 0,
+            dtype: jnp.dtype = jnp.float64,
+            _do_init: bool = True,
+            **kwargs,
+    ):
+        self.model = ViTFNQS(L_eff=config.L_eff,
+                             num_layers=config.num_layers,
+                             d_model=config.d_model,
+                             heads=config.heads,
+                             b=config.b,
+                             complex=config.complex,
+                             disorder=config.disorder,
+                             transl_invariant=config.tras_inv,
+                             two_dimensional=config.two_dim,
+        )
+        if not "return_z" in kwargs:
+            self.return_z = False
+        else:
+            self.return_z = kwargs["return_z"]
+
+        super().__init__(config, ViTFNQS, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
+
+    def __call__(self, params, spins, coups):
+        return self.model.apply(params, spins, coups, return_z=self.return_z)
+
+    def init_weights(self, rng, input_shape):
+        return self.model.init(rng, *input_shape)