Update transformer_fnqs.py

transformer_fnqs.py

@@ -99,19 +99,22 @@ 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))
 
-        amp = self.norm0(self.output_layer0(x))
+        if return_z:
+            return z
+        
+        amp = self.norm0(self.output_layer0(z))
         
         if self.complex:
-            sign = self.norm1(self.output_layer1(x))
-            z = amp + 1j*sign
+            sign = self.norm1(self.output_layer1(z))
+            out = amp + 1j*sign
         else:
-            z = amp
+            out = amp
 
-        return jnp.sum(log_cosh(z), axis=-1)
+        return jnp.sum(log_cosh(out), axis=-1)
 
 class ViTFNQS(nn.Module):
     num_layers: int
@@ -135,7 +138,7 @@ class ViTFNQS(nn.Module):
 
         self.output = OuputHead(self.d_model, complex=self.complex)
 
-    def __call__(self, spins, coups):
+    def __call__(self, spins, coups, return_z=False):
         x = jnp.atleast_2d(spins)
 
         if self.disorder:
@@ -154,6 +157,6 @@ class ViTFNQS(nn.Module):
 
         x = self.encoder(x)
 
-        z = self.output(x)
+        out = self.output(x, return_z=return_z)
 
-        return z
+        return out