NetKet 3.19.2__tar.gz → 3.20.1__tar.gz

{netket-3.19.2 → netket-3.20.1}/.gitignore

@@ -84,5 +84,12 @@ build/
 # Website #
 ###########
 docs/_build
+docs/.jupyter_cache
 _generated
 dist
+
+# Auto-generated BibTeX files #
+###############################
+references.bib
+/references.bib
+docs/references.bib

{netket-3.19.2 → netket-3.20.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.19.2
+Version: 3.20.1
 Summary: Netket : Machine Learning toolbox for many-body quantum systems.
 Project-URL: homepage, https://www.netket.org
 Project-URL: documentation, https://netket.readthedocs.io/en/latest/#
@@ -22,7 +22,7 @@ Requires-Dist: einops>=0.8.1
 Requires-Dist: equinox>=0.11.3
 Requires-Dist: flax>=0.10.2
 Requires-Dist: igraph>=0.10.0
-Requires-Dist: jax<0.6,>=0.4.35
+Requires-Dist: jax>=0.7.0
 Requires-Dist: numba>=0.59
 Requires-Dist: numpy<3.0,>=1.20
 Requires-Dist: optax<0.3,>=0.2.2
@@ -46,7 +46,7 @@ Requires-Dist: pytest-json-report>=1.3; extra == 'dev'
 Requires-Dist: pytest-xdist[psutil]>=2; extra == 'dev'
 Requires-Dist: pytest>=6; extra == 'dev'
 Requires-Dist: qutip>=5; extra == 'dev'
-Requires-Dist: ruff==0.12.2; extra == 'dev'
+Requires-Dist: ruff==0.14.1; extra == 'dev'
 Requires-Dist: wheel>=0.38.1; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: ipywidgets>=8.1.7; extra == 'docs'

netket-3.20.1/docs/sphinx_extensions/netket_citations.py

@@ -0,0 +1,256 @@
+"""
+Sphinx extension for NetKet citation system.
+
+This extension provides directives to display citation information from NetKet's
+citation system.
+"""
+
+from docutils import nodes
+from docutils.parsers.rst import Directive
+from sphinx.application import Sphinx
+from sphinx.util.docutils import SphinxDirective
+from sphinx.util import logging
+import importlib.util
+import sys
+from pathlib import Path
+
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_doi_from_bibtex(bibtex_entry):
+    """Extract DOI from bibtex entry."""
+    import re
+    
+    # Pattern to match doi field
+    doi_patterns = [
+        r'doi\s*=\s*{([^}]+)}',
+        r'doi\s*=\s*"([^"]+)"',
+        r'doi\s*=\s*([^,\s}]+)',
+    ]
+    
+    for pattern in doi_patterns:
+        match = re.search(pattern, bibtex_entry, re.IGNORECASE)
+        if match:
+            return match.group(1).strip()
+    
+    return None
+
+
+def _extract_url_from_bibtex(bibtex_entry):
+    """Extract URL from bibtex entry."""
+    import re
+    
+    # Pattern to match url field
+    url_patterns = [
+        r'url\s*=\s*{([^}]+)}',
+        r'url\s*=\s*"([^"]+)"',
+        r'url\s*=\s*([^,\s}]+)',
+    ]
+    
+    for pattern in url_patterns:
+        match = re.search(pattern, bibtex_entry, re.IGNORECASE)
+        if match:
+            return match.group(1).strip()
+    
+    return None
+
+
+def _extract_title_from_bibtex(bibtex_entry):
+    """Extract and clean title from bibtex entry."""
+    import re
+    
+    # Pattern to match title field
+    title_patterns = [
+        r'title\s*=\s*{([^}]+)}',
+        r'title\s*=\s*"([^"]+)"',
+    ]
+    
+    for pattern in title_patterns:
+        match = re.search(pattern, bibtex_entry, re.IGNORECASE)
+        if match:
+            title = match.group(1).strip()
+            # Remove LaTeX formatting
+            title = re.sub(r'\\[a-zA-Z]+{([^}]*)}', r'\1', title)
+            title = re.sub(r'[{}]', '', title)
+            return title
+    
+    return None
+
+
+def _generate_bootstrap_citations():
+    """Generate bootstrap card-based citation HTML."""
+    from netket.utils.citations import get_all_citations
+    from netket.utils.citations.core import register_citation
+    from netket import __version__ as netket_version
+    import uuid
+    
+    # Get all citations
+    citation_registry = get_all_citations()
+    
+    if not citation_registry:
+        return '<div class="alert alert-info">No citations have been registered.</div>'
+    
+    # Filter out NetKet papers
+    netket_papers = ["netket3:2022", "netket2:2019", "netket3:2021"]
+    filtered_registry = {k: v for k, v in citation_registry.items() if k not in netket_papers}
+    
+    if not filtered_registry:
+        return '<div class="alert alert-info">No algorithm-specific citations have been registered.</div>'
+    
+    # Group citations by group_id
+    groups = {}
+    for key, entry in filtered_registry.items():
+        group_id = entry.get('group_id')
+        if group_id not in groups:
+            groups[group_id] = []
+        groups[group_id].append((key, entry))
+    
+    html_parts = ['<div class="row">']
+    
+    for group_id, group_entries in groups.items():
+        # Get the first entry to extract group-level info
+        first_key, first_entry = group_entries[0]
+        condition = first_entry.get('condition', '').strip()
+        message = first_entry.get('message', 'This work used the methods described in Ref.~').strip()
+        
+        # Create bootstrap card
+        html_parts.append('<div class="col-12 mb-2">')
+        html_parts.append('<div class="card">')
+        html_parts.append('<div class="card-body p-2">')
+        
+        # Combine condition, message and references in one paragraph
+        refs_html = []
+        for ref_key, entry in group_entries:
+            bibtex = entry.get('bibtex', '')
+            if bibtex:
+                title = _extract_title_from_bibtex(bibtex)
+                doi = _extract_doi_from_bibtex(bibtex)
+                url = _extract_url_from_bibtex(bibtex)
+                
+                # Create link
+                if doi:
+                    link_url = f"https://doi.org/{doi}"
+                elif url:
+                    link_url = url
+                else:
+                    link_url = None
+                
+                if link_url and title:
+                    refs_html.append(f'(<code>{ref_key}</code>)<a href="{link_url}" target="_blank">"{title}"</a>')
+                elif title:
+                    refs_html.append(f'(<code>{ref_key}</code>)"{title}"')
+                else:
+                    refs_html.append(f'<code>{ref_key}</code>')
+            else:
+                refs_html.append(f'<code>{ref_key}</code>')
+        
+        # Create single paragraph with condition (if exists), message, and references
+        html_parts.append('<p class="card-text mb-1">')
+        if condition:
+            html_parts.append(f'<strong>{condition}:</strong> ')
+        html_parts.append(f'<em>{message}</em>')
+        if refs_html:
+            html_parts.append(' <small class="text-muted">')
+            html_parts.append(' '.join(refs_html))
+            html_parts.append('</small>')
+        html_parts.append('</p>')
+        
+        # Add collapsible BibTeX entries
+        bibtex_entries = [entry.get('bibtex', '') for _, entry in group_entries if entry.get('bibtex')]
+        if bibtex_entries:
+            html_parts.append('<details>')
+            html_parts.append('<summary><small>BibTeX entries</small></summary>')
+            html_parts.append('<pre class="mb-0"><code class="language-bibtex">')
+            html_parts.append('\n\n'.join(bibtex_entries))
+            html_parts.append('</code></pre>')
+            html_parts.append('</details>')
+        
+        html_parts.append('</div>')  # card-body
+        html_parts.append('</div>')  # card
+        html_parts.append('</div>')  # col
+    
+    html_parts.append('</div>')  # row
+    
+    return ''.join(html_parts)
+
+
+class CitationsListNode(nodes.Element):
+    """Custom node for citations list."""
+    pass
+
+
+class CitationsListDirective(SphinxDirective):
+    """
+    Directive to display all citations from NetKet's citation system.
+    
+    Example:
+        .. netket-citations::
+           :format: detailed
+    """
+    
+    has_content = False
+    optional_arguments = 0
+    option_spec = {
+        'format': str,  # 'detailed' or 'simple'
+    }
+    
+    def run(self):
+        node = CitationsListNode()
+        node['format'] = self.options.get('format', 'detailed')
+        return [node]
+
+
+def visit_citations_list_node(self, node):
+    """Visit function for CitationsListNode."""
+    format_type = node['format']
+    
+    try:
+        # Import NetKet to load citation registrations
+        import netket as nk
+        from netket.utils.citations import get_all_citations
+        import re
+        import uuid
+        
+        if format_type == 'detailed':
+            citations_html = _generate_bootstrap_citations()
+        else:
+            from netket.utils.citations.docs import get_citations_for_docs
+            citations_html = get_citations_for_docs(format_type)
+        
+        self.body.append(citations_html)
+        
+    except ImportError as e:
+        logger.warning(f"Could not import NetKet for citations: {e}")
+        self.body.append('<div class="admonition warning"><p class="admonition-title">Warning</p>')
+        self.body.append('<p>Citations could not be loaded. NetKet import failed.</p></div>')
+    except Exception as e:
+        logger.warning(f"Error generating citations: {e}")
+        self.body.append('<div class="admonition error"><p class="admonition-title">Error</p>')
+        self.body.append(f'<p>Error generating citations: {e}</p></div>')
+
+
+def depart_citations_list_node(self, node):
+    """Depart function for CitationsListNode."""
+    pass
+
+
+def setup(app: Sphinx):
+    """Setup function for the Sphinx extension."""
+    
+    # Add the custom directive
+    app.add_directive('netket-citations', CitationsListDirective)
+    
+    # Add the custom node and its handlers
+    app.add_node(
+        CitationsListNode,
+        html=(visit_citations_list_node, depart_citations_list_node),
+        latex=(visit_citations_list_node, depart_citations_list_node),
+        text=(visit_citations_list_node, depart_citations_list_node),
+    )
+    
+    return {
+        'version': '1.0',
+        'parallel_read_safe': True,
+        'parallel_write_safe': True,
+    }

{netket-3.19.2 → netket-3.20.1}/netket/__init__.py

@@ -12,18 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ._version import version as __version__  # noqa: F401
-
-from .utils import config
-
-from . import utils
-from . import errors
-
 __all__ = [
+    "callbacks",
+    "errors",
     "exact",
+    "experimental",
     "graph",
-    "callbacks",
     "hilbert",
+    "logging",
+    "models",
+    "nn",
     "operator",
     "optimizer",
     "sampler",
@@ -31,8 +29,17 @@ __all__ = [
     "utils",
     "vqs",
     "nn",
+    "symmetry",
+    "cite",
+    "config",
 ]
 
+from ._version import version as __version__  # noqa: F401
+
+from .utils import config
+
+from . import utils
+from . import errors
 
 from . import jax
 from . import stats
@@ -51,9 +58,21 @@ from . import (
     sampler,
     vqs,
     optimizer,
+    symmetry,
 )
 
+from . import experimental
+
 
 # Main applications
 from .driver import VMC
 from .driver import SteadyState
+
+# Citation system
+from netket.utils.citations import cite
+
+# Show tips if in interactive mode
+from netket.utils._tips import show_random_tip
+
+show_random_tip()
+del show_random_tip

{netket-3.19.2/netket/experimental → netket-3.20.1/netket/_src}/driver/vmc_sr.py

@@ -1,13 +1,17 @@
 from typing import Any
 from collections.abc import Callable
+from warnings import warn
 
 import jax
 from jax.flatten_util import ravel_pytree
 
 from netket import jax as nkjax
+from netket.hilbert import SpinOrbitalFermions
 from netket.optimizer.solver import cholesky
 from netket.vqs import MCState, FullSumState
 from netket.utils import timing, struct
+from netket.utils.citations import reference
+from netket.utils.deprecation import DeprecatedArg
 from netket.utils.types import ScalarOrSchedule, Optimizer, Array, PyTree
 from netket.jax._jacobian.default_mode import JacobianMode
 from netket.operator import AbstractOperator
@@ -37,22 +41,33 @@ def VMC_SRt(
     optimizer: Optimizer,
     *,
     diag_shift: ScalarOrSchedule,
-    linear_solver_fn: Callable[[jax.Array, jax.Array], jax.Array] = cholesky,
+    linear_solver: Callable[[jax.Array, jax.Array], jax.Array] = cholesky,
     mode: str | None = None,
-    jacobian_mode: str | None = None,
-    variational_state: MCState = None,
+    jacobian_mode: str | None | DeprecatedArg = DeprecatedArg(),
+    variational_state: MCState,
+    linear_solver_fn: (
+        Callable[[jax.Array, jax.Array], jax.Array] | DeprecatedArg
+    ) = DeprecatedArg(),
 ):
-    if mode is None:
-        mode = jacobian_mode
-    elif mode is not None and jacobian_mode is not None:
-        raise ValueError(
-            "`jacobian_mode` is deprecated and renamed to `mode`. Just declare `mode`."
+    if not isinstance(jacobian_mode, DeprecatedArg):
+        if mode is not None:
+            raise ValueError(
+                "`jacobian_mode` is deprecated and renamed to `mode`. Just declare `mode`."
+            )
+        warn(
+            """
+            The keyword argument `linear_solver_fn` is deprecated in favour of `linear_solver`.
+            """,
+            category=FutureWarning,
+            stacklevel=2,
         )
+        mode = jacobian_mode
 
     return VMC_SR(
         hamiltonian,
         optimizer,
         diag_shift=diag_shift,
+        linear_solver=linear_solver,
         linear_solver_fn=linear_solver_fn,
         mode=mode,
         variational_state=variational_state,
@@ -61,6 +76,16 @@ def VMC_SRt(
     )
 
 
+@reference(
+    "Goldshlager2023Spring",
+    condition="If using VMC_SR with momentum != 0",
+    message="This work used the SPRING optimization algorithm described in Ref.",
+)
+@reference(
+    ["Chen2024minsr", "Rende2024minsr"],
+    condition="If using minSR, VMC_SR(use_ntk=True) ",
+    message="This work used the efficient kernel trick for SR described in Refs.",
+)
 class VMC_SR(AbstractVariationalDriver):
     r"""
     Energy minimization using Variational Monte Carlo (VMC) and **Stochastic Reconfiguration/Natural Gradient Descent**.
@@ -94,7 +119,7 @@ class VMC_SR(AbstractVariationalDriver):
 
     .. code-block:: python
 
-        linear_solver_fn(A: Matrix, b: vector) -> tuple[jax.Array[vector], dict]
+        linear_solver(A: Matrix, b: vector) -> tuple[jax.Array[vector], dict]
 
     Where the vector is the solution and the dictionary may contain additional information about the solver or be None.
     The standard solver is based on the Cholesky decomposition :func:`~netket.optimizer.solver.cholesky`, but any other
@@ -208,7 +233,7 @@ class VMC_SR(AbstractVariationalDriver):
     _chunk_size_bwd: int | None = struct.field(serialize=False)
     _use_ntk: bool = struct.field(serialize=False)
     _on_the_fly: bool = struct.field(serialize=False)
-    _linear_solver_fn: Any = struct.field(serialize=False)
+    _linear_solver: Any = struct.field(serialize=False)
 
     # Internal things cached
     _unravel_params_fn: Any = struct.field(serialize=False)
@@ -229,8 +254,11 @@ class VMC_SR(AbstractVariationalDriver):
         diag_shift: ScalarOrSchedule,
         proj_reg: ScalarOrSchedule | None = None,
         momentum: ScalarOrSchedule | None = None,
-        linear_solver_fn: Callable[[Array, Array], Array] = cholesky,
-        variational_state: MCState = None,
+        linear_solver: Callable[[Array, Array], Array] = cholesky,
+        linear_solver_fn: (
+            Callable[[Array, Array], Array] | DeprecatedArg
+        ) = DeprecatedArg(),
+        variational_state: MCState,
         chunk_size_bwd: int | None = None,
         mode: JacobianMode | None = None,
         use_ntk: bool | None = None,
@@ -247,7 +275,7 @@ class VMC_SR(AbstractVariationalDriver):
         Args:
             hamiltonian: The Hamiltonian of which the ground-state is to be found.
             optimizer: The optimizer to use for the parameter updates. To perform proper
-                SR/NGD optimization this should be an instance of `optax.sgd`, but can be
+                SR/NGD optimization this should be an instance of :func:`optax.sgd`, but can be
                 any other optimizer if you are brave.
             variational_state: The variational state to optimize.
             diag_shift: The diagonal regularization parameter :math:`\lambda` for the QGT/NTK.
@@ -261,11 +289,32 @@ class VMC_SR(AbstractVariationalDriver):
                 Thus the  amplification is at most a factor of :math:`A(0.9)=2.3` or
                 :math:`A(0.99)=7.1`. Values around ``momentum = 0.8`` empirically work well.
                 (Defaults to None)
-            linear_solver_fn: The linear solver function to use for the NGD solver.
+            linear_solver: The linear solver function to use for the NGD solver. Defaults to
+                :func:`netket.optimizer.solver.cholesky`, but can use other solvers from there. In general:
+
+                - :func:`~netket.optimizer.solver.cholesky` is faster because it relies on LU decomposition
+                  instead of a full diagonalization, but it is more prone to numerical instabilities,
+                  especially with explicitly simmetrized networks or very singular QGT/NTKs. Often the issue
+                  is that your matrix is numerically not hermitian/positive semidefinite because of numerical
+                  errors, and this breaks the method. **If you see NaNs in your weights
+                  during your optimization, this is likely the cause.**
+                - :func:`~netket.optimizer.solver.pinv_smooth` (a smoothed variant of
+                  :func:`~netket.optimizer.solver.pinv`) is considerably more stable and usually does
+                  not cause NaNs, but it is also considerably more expensive.
+                - :func:`~jax.scipy.sparse.linalg.cg` and other iterative solvers can sometimes be faster, but
+                  the quality of the solution is bad and they can lead to unpredictable step times because
+                  the number of CG iterations might vary with the condition number. While we used those a
+                  lot in the past, there is considerable evidence that they should be avoided (see Chen &
+                  Heyl Nature Physics).
+                In general, if you are not bottlenecked by the linear solver, it is a good idea to use the
+                more reliable :func:`~netket.optimizer.solver.pinv_smooth`.
             mode: The mode used to compute the jacobian of the variational state.
-                Can be `'real'` or `'complex'`. Real can be used for real-valued wavefunctions
-                with a sign, to truncate the arbitrary phase of the wavefunction. This leads
-                to lower computational cost.
+                Can be ``'real'`` or ``'complex'``. Real can be used for real-valued wavefunctions
+                with a sign, to truncate the arbitrary phase of the wavefunction. In complex mode, the QGT/NTK
+                is concretized as a real-valued :math:`2N \times 2N` where :math:`N` is either the number of
+                parameters or number of samples.
+                If your wavefunctino is real (as is usually the case for fermionic hamiltonians) you should
+                really set this to `real`.
             on_the_fly: Whether to compute the QGT or NTK using lazy evaluation methods.
                 This usually requires less memory. (Defaults to None, which will
                 automatically chose the potentially best method).
@@ -276,6 +325,17 @@ class VMC_SR(AbstractVariationalDriver):
                 SR and minSR. (Defaults to None, which will automatically choose the best
                 method)
         """
+        # TODO: Deprecated in September 2025, netket 3.20
+        if not isinstance(linear_solver_fn, DeprecatedArg):
+            warn(
+                """
+                The keyword argument `linear_solver_fn` is deprecated in favour of `linear_solver`.
+                """,
+                category=FutureWarning,
+                stacklevel=2,
+            )
+            linear_solver = linear_solver_fn
+
         if isinstance(variational_state, FullSumState):
             raise TypeError(
                 "NGD drivers do not support FullSumState. Please use 'standard' drivers with SR."
@@ -294,13 +354,30 @@ class VMC_SR(AbstractVariationalDriver):
         elif on_the_fly and not use_ntk:
             raise ValueError(
                 """
-                `onthefly` is only supported when `use_ntk=True`.
+                `on_the_fly=True` is only supported when `use_ntk=True`.
 
                 We plan to support this mode for the standard NGD in the future.
-                In the meantime, use a standard VMC+SR QGTOnTheFly.
+                In the meantime, use a standard VMC+SR with QGTOnTheFly if you
+                really want to use `on_the_fly=True`
                 """
             )
 
+        # for most fermionic calculations you want to use real mode, but the auto
+        # detect will pick complex
+        if mode is None and isinstance(variational_state.hilbert, SpinOrbitalFermions):
+            warn(
+                "`mode` not selected when working with Fermionic systems. Automatically falling back "
+                "to `mode='complex'` but for most fermionic calculations where the wave-function is real "
+                "valued with a sign (not a complex phase) `mode='real'` will truncate the imaginary part "
+                "of the QGT and lead to the same optimization but at a considerably lower cost (~8 times less)."
+                "\n\nIT'S LIKELY YOU WANT TO SPECIFY `mode='real'`.\n\n"
+            )
+
+        # Default chunk_size_bwd to the variational's state chunk size. In any case,
+        # it usually should be smaller than the chunk size of variational state.
+        if chunk_size_bwd is None:
+            chunk_size_bwd = variational_state.chunk_size
+
         self._ham = hamiltonian
 
         self.diag_shift = diag_shift
@@ -312,7 +389,7 @@ class VMC_SR(AbstractVariationalDriver):
         self.mode = mode
         self._on_the_fly = on_the_fly
 
-        self._linear_solver_fn = linear_solver_fn
+        self._linear_solver = linear_solver
 
         _, unravel_params_fn = ravel_pytree(self.state.parameters)
         self._unravel_params_fn = jax.jit(unravel_params_fn)
@@ -370,7 +447,7 @@ class VMC_SR(AbstractVariationalDriver):
             self.state.model_state,
             samples,
             diag_shift=diag_shift,
-            solver_fn=self._linear_solver_fn,
+            solver_fn=self._linear_solver,
             mode=self.mode,
             proj_reg=proj_reg,
             momentum=momentum,
@@ -417,9 +494,6 @@ class VMC_SR(AbstractVariationalDriver):
 
     @mode.setter
     def mode(self, mode: str | JacobianMode | None):
-        # TODO: Add support for 'onthefly' mode
-        # At the moment, onthefly is only supported for use_ntk=True.
-        # We raise a warning if the user tries to use it with use_ntk=False.
         if mode is None:
             mode = nkjax.jacobian_default_mode(
                 self.state._apply_fun,
@@ -444,7 +518,9 @@ class VMC_SR(AbstractVariationalDriver):
     @property
     def on_the_fly(self) -> bool:
         """
-        Whether
+        Whether to use a lazy implementation of th NTK or QGT which does not concretize the jacobian.
+
+        This usually requires less memory.
         """
         return self._on_the_fly
 
@@ -469,5 +545,7 @@ class VMC_SR(AbstractVariationalDriver):
     @chunk_size_bwd.setter
     def chunk_size_bwd(self, value: int | None):
         if not isinstance(value, int | None):
-            raise TypeError("chunk_size must be an integer or None")
+            raise TypeError("chunk_size must be an integer or None.")
+        elif isinstance(value, int) and value <= 0:
+            raise ValueError("chunk_size_bwd must be a positive integer.")
         self._chunk_size_bwd = value

{netket-3.19.2 → netket-3.20.1}/netket/_src/ngd/sr.py

@@ -3,11 +3,13 @@ from functools import partial
 
 import jax
 import jax.numpy as jnp
-from jax.sharding import PositionalSharding
+from jax.sharding import NamedSharding, PartitionSpec as P
 
 from netket import jax as nkjax
 from netket.utils.types import Array
 
+from netket import config
+
 from netket._src.ngd.kwargs import ensure_accepts_kwargs
 
 
@@ -34,12 +36,16 @@ def _compute_sr_update(
     # Typically solvers only accept the matrix and the right-hand side.
     solver_fn = ensure_accepts_kwargs(solver_fn, "dv")
 
-    if (momentum is not None) or (old_updates is not None) or (proj_reg is not None):
-        raise ValueError("Not implemented")
+    if proj_reg is not None:
+        raise ValueError("proj_reg not implemented for SR")
 
     # (np, #ns) x (#ns) -> (np) - where the sum over #ns is done automatically
     F = O_L.T @ dv
 
+    # Add momentum term to F: F + λ μ dθ_{i-1}
+    if momentum is not None:
+        F = F + diag_shift * momentum * old_updates
+
     # This does the contraction (np, #ns) x (#ns, np) -> (np, np).
     matrix = O_L.T @ O_L
     matrix_side = matrix.shape[-1]  # * it can be ns or 2*ns, depending on mode
@@ -57,14 +63,18 @@ def _compute_sr_update(
     else:
         info = {}
 
+    if momentum is not None:
+        old_updates = updates
+
     # If complex mode and we have complex parameters, we need
     # To repack the real coefficients in order to get complex updates
     if mode == "complex" and nkjax.tree_leaf_iscomplex(params_structure):
         num_p = updates.shape[-1] // 2
         updates = updates[:num_p] + 1j * updates[num_p:]
 
-    updates = jax.lax.with_sharding_constraint(
-        updates, PositionalSharding(jax.devices()).replicate()
-    )
+    if config.netket_experimental_sharding:
+        updates = jax.lax.with_sharding_constraint(
+            updates, NamedSharding(jax.sharding.get_abstract_mesh(), P())
+        )
 
     return updates, old_updates, info

{netket-3.19.2 → netket-3.20.1}/netket/_src/ngd/sr_srt_common.py

@@ -101,7 +101,7 @@ def _sr_srt_common(
         diag_shift: The diagonal shift of the stochastic reconfiguration matrix. Typical values are 1e-4 ÷ 1e-3. Can also be an optax schedule.
         proj_reg: Weight before the matrix `1/N_samples \\bm{1} \\bm{1}^T` used to regularize the linear solver in SPRING.
         momentum: Momentum used to accumulate updates in SPRING.
-        linear_solver_fn: Callable to solve the linear problem associated to the updates of the parameters.
+        linear_solver: Callable to solve the linear problem associated to the updates of the parameters.
         mode: The mode used to compute the jacobian of the variational state. Can be `'real'` or `'complex'` (defaults to the dtype of the output of the model).
 
     Returns: