how to implement batch-computation via jnp.vmap by myself

[miaoL1]

I am constructing the log of the overlap between the variant state and the basis.
From the teaching documents, it can be seen that most models utilize JNP's built-in functions to automatically handle the calculation of batch dimensions. For example, the meanfield ansatz of the Ising model:

class MF(nn.Module):
    @nn.compact
    def __call__(self, x):
        lam = self.param(
            "lambda", nn.initializers.normal(), (1,), float
        )
        p = nn.log_sigmoid(lam*x)
        return 0.5 * jnp.sum(p, axis=-1)

However, when I tried to construct a model by myself, I encountered bugs. (I am a beginner using jax, so what I may encounter is probably just coding issues with jax)

class amplitude(nn.Module):
  hilbert: nkx.hilbert.SpinOrbitalFermions
  graph: nk.graph.Graph
  t: float = 1
  def compute_energy(mu0,mu1,mu2 x):
      N = jnp.int64(x.shape[-1] // 2)
      energy = 0.0
      for i in range(int(N // 2)):
          idx = abs(x[i] - x[i + int(N // 2)] + x[i + N] - x[i + N + int(N // 2)])
          energy += jnp.array([mu0,mu1,mu2])[idx]
      return energy
  @nn.compact
  def __call__(self,x):
    lamda = self.param("lamda",nn.initializers.normal(),(1,),jnp.float64)
    mu0 = self.param("mu0",nn.initializers.normal(),(1,),jnp.float64)
    mu1 = self.param("mu1",nn.initializers.normal(),(1,),jnp.float64)
    mu2 = self.param("mu2",nn.initializers.normal(),(1,),jnp.float64)
    energy = jax.vmap(self.compute_energy, in_axes=(None, None,None,0))(mu0,mu1,mu2, x)
    return jnp.array(energy)

As shown above, I tried to compute x with batch size, but Type error emerged, saying that "compute_energy required 4 parameters but 5 are given". I guess this might have something to do with defining the function within the class.
I'd appreciate it if anyone can help me correct the bug.

(BTW, I code with Torch a lot, and I'm very used to unsqueezing the data to get the batch dimension, then calculating with broadcast mechanism conveniently. Is there any similar method to do so in Jax?)

[PhilipVinc]

I think this is a flax-specific question and not really connected to netket. Ask it on their forum!

I think it does not work because of the way flax handles functions in those modules. Wrapping compute_energy with @nn.nowrap might solve your problem.