more docstring fixes

netket/sampler/metropolis.py

@@ -33,17 +33,6 @@ class MetropolisRule:
     """
     Base class for Transition rules of Metropolis, such as Local, Exchange, Hamiltonian
     and several others.
-
-    Methods:
-        - init_state(rule, sampler, machine, params), constructing the state of the rule.
-        By default the state is None.
-
-        - reset(rule, sampler, machine, parameters, state), resets the state of the rule,
-        by default returns None.
-        - transition(rule, sampler, machine, parameters, state, key, σ), which returns a
-        new configuration given the configuration σ and rng key key.
-        - random_state(rule, sampler, machine, parameters, state, key), which returns a
-        new random configuration. By default this uses standard hilbert random_state.
     """
 
     def init_state(
@@ -55,12 +44,15 @@ class MetropolisRule:
     ) -> Optional[Any]:
         """
         Initialises the optional internal state of the Metropolis Sampler Transition
-        Rule. The provided key is unique and does not need to be splitted
+        Rule.
+
+        The provided key is unique and does not need to be splitted.
         It should return an immutable datastructure.
 
         Arguments:
             sampler: The Metropolis sampler
             machine: The forward evaluation function of the model, accepting PyTrees of parameters and inputs.
+            params: The dict of variables needed to evaluate the model.
             key: A Jax PRNGKey rng state.
 
         Returns:
@@ -81,7 +73,7 @@ class MetropolisRule:
         Arguments:
             sampler: The Metropolis sampler
             machine: The forward evaluation function of the model, accepting PyTrees of parameters and inputs.
-            key: A Jax PRNGKey rng state.
+            params: The dict of variables needed to evaluate the model.
             sampler_state: The current state of the sampler. Should not modify it.
 
         Returns:
@@ -109,6 +101,18 @@ class MetropolisRule:
         state: SamplerState,
         key: PRNGKey,
     ):
+        """
+        Generates a random state compatible with this rule.
+
+        By default this calls :func:`netket.hilbert.random.random_state`.
+
+        Arguments:
+            sampler: the sampler
+            machine: the function to evaluate the model
+            parameters: the parameters of the model
+            state: the current sampler state
+            key: the PRNGKey to use to generate the random state
+        """
         return sampler.hilbert.random_state(
             key, size=sampler.n_batches, dtype=sampler.dtype
         )

netket/sampler/rules/exchange.py

@@ -109,24 +109,24 @@ def ExchangeRule(
     possible couples (clusters).
 
     This rule acts on two local degree of freedom :math:`s_i` and :math:`s_j`,
-    and proposes a new state: :math:`s_1 \dots s^\prime_i \dots s^\prime_j \dots s_N`,
-    where in general :math:`s^\prime_i \neq s_i` and :math:`s^\prime_j \neq s_j`.
+    and proposes a new state: :math:`s_1 \\dots s^\\prime_i \\dots s^\\prime_j \\dots s_N`,
+    where in general :math:`s^\\prime_i \\neq s_i` and :math:`s^\\prime_j \\neq s_j`.
     The sites :math:`i` and :math:`j` are also chosen to be within a maximum graph
-    distance of :math:`d_{\mathrm{max}}`.
+    distance of :math:`d_{\\mathrm{max}}`.
 
     The transition probability associated to this sampler can
     be decomposed into two steps:
 
-    1. A pair of indices :math:`i,j = 1\dots N`, and such
-       that :math:`\mathrm{dist}(i,j) \leq d_{\mathrm{max}}`,
+    1. A pair of indices :math:`i,j = 1\\dots N`, and such
+       that :math:`\\mathrm{dist}(i,j) \\leq d_{\\mathrm{max}}`,
        is chosen with uniform probability.
-    2. The sites are exchanged, i.e. :math:`s^\prime_i = s_j` and :math:`s^\prime_j = s_i`.
+    2. The sites are exchanged, i.e. :math:`s^\\prime_i = s_j` and :math:`s^\\prime_j = s_i`.
 
     Notice that this sampling method generates random permutations of the quantum
     numbers, thus global quantities such as the sum of the local quantum numbers
     are conserved during the sampling.
     This scheme should be used then only when sampling in a
-    region where :math:`\sum_i s_i = \mathrm{constant}` is needed,
+    region where :math:`\\sum_i s_i = \\mathrm{constant}` is needed,
     otherwise the sampling would be strongly not ergodic.
 
     Args: