Sharding errors in sampler when using `reset_chains` in multi-node setting about netket

[PhilipVinc]

Can you try to add a @jax.jit to the MetropolisSampler.reset method?

Otherwise @inailuig might know better.

[inailuig]

#1721 will fix this

The issue is that the key in

netket/netket/sampler/metropolis.py

Line 396 in f378bd4

σ = sampler.rule.random_state(sampler, machine, parameters, state, rng)

is sharded after the first call, but random_state is not jitted, so it errors if its called with non fully addressable arrays.

[inailuig]

@PhilipVinc mabye we can do a bugfix release which includes #1727
and I'll complete this todo

netket/netket/experimental/sampler/metropolis_pt.py

Lines 160 to 165 in f378bd4

	σ = sampler.rule.random_state(sampler, machine, params, rule_state, rng)
	if config.netket_experimental_sharding and jax.device_count() > 1:
	# TODO If we end up rewriting the hilbert spaces in jax then we can avoid
	# this and instead jit the random state with the correct out_shardings
	σ = distribute_to_devices_along_axis(σ, axis=0)

in the other PR to solve it properly later.

[PhilipVinc]

Solved by #1727 , released few moments ago in netket 3.11.1

[therooler]

Thanks guys, I appreciate the quick fix!