import netket as nk
import jax
import jax.numpy as jnp

g1 = nk.graph.Chain(4)
g2 = nk.graph.Chain(1, pbc=False)
hi1 = nk.hilbert.Fock(n_max=8, N=g1.n_nodes)
hi2 = nk.hilbert.Spin(s=1/2, N=g2.n_nodes)
hi = hi1*hi2

k = jax.random.key(123)
x = hi.random_state(k, (23,))
print(jax.tree.map(lambda x: x.shape, x))

ma1 = nk.models.RBM
ma2 = nk.models.RBM

from netket.utils.samples_pytree import SampleWrapperExample
from flax import linen as nn
class ProdModel(nn.Module):
    models: tuple[nn.Module]
    @nn.compact
    def __call__(self, x):
        assert isinstance (x, SampleWrapperExample)
        assert len(x.sub_states) == len(self.models)
        return sum(mi()(xi) for mi, xi in zip(self.models, x.sub_states))

ma = ProdModel((ma1, ma2))

sa = nk.sampler.MetropolisLocal(hi)
vs = nk.vqs.MCState(sa, ma)
vs.sample()

S = vs.quantum_geometric_tensor()
S@vs.parameters