NetKet 3.20.2__tar.gz → 3.20.3__tar.gz

{netket-3.20.2 → netket-3.20.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.20.2
+Version: 3.20.3
 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/#

{netket-3.20.2 → netket-3.20.3}/netket/_src/ngd/srt.py

@@ -34,13 +34,16 @@ def _compute_srt_update(
     if momentum is not None:
         dv -= momentum * (O_L @ old_updates)
 
+    # Here we reshard the jacobian from being sharded along the sample axis to being
+    # sharded along the parameter axis, which we pad
     # (#ns, np) -> (ns, #np)
+    # In theory jax could figure this out, but maybe he does not, so we do it by hand.
     O_LT = O_L
     if config.netket_experimental_sharding:
-        nkjax.sharding.pad_axis_for_sharding(O_LT, axis=1, padding_value=0.0)
+        O_LT = nkjax.sharding.pad_axis_for_sharding(O_LT, axis=1, padding_value=0.0)
         O_LT = jax.lax.with_sharding_constraint(
             O_LT,
-            NamedSharding(jax.sharding.get_abstract_mesh(), P("S", None)),
+            NamedSharding(jax.sharding.get_abstract_mesh(), P(None, "S")),
         )
         dv = jax.lax.with_sharding_constraint(
             dv, NamedSharding(jax.sharding.get_abstract_mesh(), P())

{netket-3.20.2 → netket-3.20.3}/netket/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.20.2'
-__version_tuple__ = version_tuple = (3, 20, 2)
+__version__ = version = '3.20.3'
+__version_tuple__ = version_tuple = (3, 20, 3)
 
 __commit_id__ = commit_id = None