Upload model

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:73bcef74adf67486945e05ed20e67eb48d071fbe22c8b3de8aed501b2f417df7
+oid sha256:1e327504c22eaeca2ba25d074c771e63462500f1301ee258b9afef233455a82a
 size 3490136

transformer.py

@@ -101,12 +101,14 @@ class OuputHead(nn.Module):
         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):
+    def __call__(self, x, return_z=False):
 
-        x = self.out_layer_norm(x.sum(axis=1))
+        z = self.out_layer_norm(x.sum(axis=1))
+        if return_z:
+            return z
 
-        amp = self.norm2(self.output_layer0(x))
-        sign = self.norm3(self.output_layer1(x))
+        amp = self.norm2(self.output_layer0(z))
+        sign = self.norm3(self.output_layer1(z))
         
         z = amp + 1j*sign
 
@@ -129,13 +131,13 @@ class ViT(nn.Module):
         self.output = OuputHead(self.d_model)
 
 
-    def __call__(self, spins):
+    def __call__(self, spins, return_z=False):
         x = jnp.atleast_2d(spins)
 
         x = self.patches_and_embed(x)
 
         x = self.encoder(x)
 
-        z = self.output(x)
+        z = self.output(x, return_z=return_z)
 
         return z

vitnqs_model.py

@@ -24,11 +24,15 @@ class ViTNQSModel(FlaxPreTrainedModel):
                          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, ViT, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
 
     def __call__(self, params, spins):
-        return self.model.apply(params, spins)
+        return self.model.apply(params, spins, self.return_z)
 
     def init_weights(self, rng, input_shape):
         return self.model.init(rng, input_shape)