NetKet 3.22.3__tar.gz → 3.23.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (444) hide show
  1. {netket-3.22.3 → netket-3.23.0}/.gitignore +3 -0
  2. {netket-3.22.3 → netket-3.23.0}/PKG-INFO +4 -3
  3. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/base.py +20 -1
  4. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/convergence_stopping.py +33 -9
  5. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/early_stopping.py +24 -5
  6. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/invalid_loss_stopping.py +22 -3
  7. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/progressbar.py +31 -16
  8. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/save_state.py +36 -6
  9. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/stop_at_time.py +10 -1
  10. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/timeout.py +35 -2
  11. {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/abstract_dynamics_driver.py +1 -2
  12. {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/abstract_optimization_driver.py +4 -3
  13. {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/abstract_variational_driver.py +36 -12
  14. {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/infidelity_sr.py +7 -2
  15. {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/vmc_sr.py +56 -9
  16. {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/hdf5_log.py +17 -12
  17. {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/mlflow_log.py +26 -5
  18. {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/tensorboard.py +15 -3
  19. netket-3.23.0/netket/_src/logging/wandb_log.py +376 -0
  20. {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/srt_onthefly.py +53 -8
  21. {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/common.py +25 -5
  22. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/pyscf_api.py +1 -1
  23. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/pyscf_utils.py +1 -1
  24. netket-3.23.0/netket/_src/solvers/jaxmg_interface.py +226 -0
  25. {netket-3.22.3 → netket-3.23.0}/netket/_src/solvers/nan_fallback.py +8 -0
  26. {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/operations.py +1 -1
  27. {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/representation.py +1 -1
  28. netket-3.23.0/netket/_src/vqs/expect_to_precision.py +221 -0
  29. {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/freeze_vstate.py +16 -50
  30. {netket-3.22.3 → netket-3.23.0}/netket/_version.py +2 -2
  31. {netket-3.22.3 → netket-3.23.0}/netket/callbacks.py +1 -0
  32. {netket-3.22.3 → netket-3.23.0}/netket/driver/steady_state.py +4 -6
  33. {netket-3.22.3 → netket-3.23.0}/netket/driver/vmc.py +4 -6
  34. {netket-3.22.3 → netket-3.23.0}/netket/experimental/driver/tdvp.py +13 -0
  35. netket-3.23.0/netket/experimental/driver/tdvp_common.py +518 -0
  36. {netket-3.22.3 → netket-3.23.0}/netket/experimental/driver/tdvp_schmitt.py +17 -1
  37. {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_integrator_state.py +5 -2
  38. {netket-3.22.3 → netket-3.23.0}/netket/graph/graph.py +1 -1
  39. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/tensor_hilbert_discrete.py +1 -1
  40. {netket-3.22.3 → netket-3.23.0}/netket/jax/__init__.py +1 -0
  41. netket-3.23.0/netket/jax/_apply_packed.py +461 -0
  42. {netket-3.22.3 → netket-3.23.0}/netket/jax/_ntk/logic.py +18 -3
  43. {netket-3.22.3 → netket-3.23.0}/netket/jax/_utils_tree.py +16 -0
  44. {netket-3.22.3 → netket-3.23.0}/netket/logging.py +1 -0
  45. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/solver/solvers.py +116 -64
  46. {netket-3.22.3 → netket-3.23.0}/netket/sampler/metropolis_numpy.py +13 -2
  47. {netket-3.22.3 → netket-3.23.0}/netket/sampler/parallel_tempering.py +8 -2
  48. {netket-3.22.3 → netket-3.23.0}/netket/tools/djaxrun.py +28 -0
  49. {netket-3.22.3 → netket-3.23.0}/netket/utils/_tips.py +1 -1
  50. {netket-3.22.3 → netket-3.23.0}/netket/utils/array.py +1 -1
  51. {netket-3.22.3 → netket-3.23.0}/netket/utils/config_flags.py +24 -4
  52. {netket-3.22.3 → netket-3.23.0}/netket/utils/history/history.py +24 -1
  53. netket-3.23.0/netket/utils/optional_deps.py +81 -0
  54. {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/fields.py +5 -2
  55. {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/pytree.py +24 -0
  56. {netket-3.22.3 → netket-3.23.0}/netket/vqs/base.py +26 -4
  57. {netket-3.22.3 → netket-3.23.0}/netket/vqs/full_summ/state.py +11 -0
  58. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/common.py +13 -4
  59. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/kernels.py +1 -1
  60. netket-3.23.0/netket/vqs/mc/kernels_packed.py +212 -0
  61. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect.py +4 -1
  62. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_chunked.py +12 -3
  63. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/state.py +62 -14
  64. {netket-3.22.3 → netket-3.23.0}/pyproject.toml +3 -3
  65. netket-3.22.3/netket/_src/driver/abstract_variational_driver_old.py +0 -443
  66. netket-3.22.3/netket/_src/vqs/expect_to_precision.py +0 -164
  67. netket-3.22.3/netket/experimental/driver/tdvp_common.py +0 -457
  68. netket-3.22.3/netket/utils/optional_deps.py +0 -39
  69. {netket-3.22.3 → netket-3.23.0}/LICENSE +0 -0
  70. {netket-3.22.3 → netket-3.23.0}/README.md +0 -0
  71. {netket-3.22.3 → netket-3.23.0}/docs/sphinx_extensions/netket_citations.py +0 -0
  72. {netket-3.22.3 → netket-3.23.0}/netket/__init__.py +0 -0
  73. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/auto_chunk_size.py +0 -0
  74. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/auto_slurm_requeue.py +0 -0
  75. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/callback_list.py +0 -0
  76. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/legacy_wrappers.py +0 -0
  77. {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/observables.py +0 -0
  78. {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/base.py +0 -0
  79. {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/json_log.py +0 -0
  80. {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/runtime_log.py +0 -0
  81. {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/state_log.py +0 -0
  82. {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/kwargs.py +0 -0
  83. {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/sr.py +0 -0
  84. {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/sr_srt_common.py +0 -0
  85. {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/srt.py +0 -0
  86. {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/apply_operator/functional.py +0 -0
  87. {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/apply_operator/linen.py +0 -0
  88. {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/apply_operator/nnx.py +0 -0
  89. {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/api.py +0 -0
  90. {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/functional.py +0 -0
  91. {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/linen.py +0 -0
  92. {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/nnx.py +0 -0
  93. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/correlator/__init__.py +0 -0
  94. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/correlator/correlator.py +0 -0
  95. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/correlator/expect.py +0 -0
  96. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/infidelity/__init__.py +0 -0
  97. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/infidelity/exact.py +0 -0
  98. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/infidelity/expect.py +0 -0
  99. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/infidelity/infidelity_operator.py +0 -0
  100. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/renyi2/S2_operator.py +0 -0
  101. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/renyi2/__init__.py +0 -0
  102. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/renyi2/exact.py +0 -0
  103. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/renyi2/expect.py +0 -0
  104. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/variance/__init__.py +0 -0
  105. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/variance/exact.py +0 -0
  106. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/variance/expect.py +0 -0
  107. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/variance/variance_operator.py +0 -0
  108. {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/vscore.py +0 -0
  109. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/normal_order_utils.py +0 -0
  110. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/conversion.py +0 -0
  111. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/fermihubbard.py +0 -0
  112. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/kernels.py +0 -0
  113. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/operator_data.py +0 -0
  114. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/operators.py +0 -0
  115. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/base.py +0 -0
  116. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/construct.py +0 -0
  117. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/fermion.py +0 -0
  118. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/spin.py +0 -0
  119. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/trace_utils.py +0 -0
  120. {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/spin_flip.py +0 -0
  121. {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/local_estimators.py +0 -0
  122. {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/__init__.py +0 -0
  123. {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/accumulator.py +0 -0
  124. {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/accumulator_batch.py +0 -0
  125. {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/kernels.py +0 -0
  126. {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/canonical_representation.py +0 -0
  127. {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/labeled_representation.py +0 -0
  128. {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/labeled_representation_coset_filter.py +0 -0
  129. {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/representation_construction.py +0 -0
  130. {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/spin_flip_representation.py +0 -0
  131. {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/translation_coset_filter.py +0 -0
  132. {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/translation_representation.py +0 -0
  133. {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/check_mc_convergence.py +0 -0
  134. {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/fermion_mf/__init__.py +0 -0
  135. {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/fermion_mf/expectation.py +0 -0
  136. {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/fermion_mf/qgt.py +0 -0
  137. {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/fermion_mf/state.py +0 -0
  138. {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/plot_mc_convergence.py +0 -0
  139. {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/transformed_vstate.py +0 -0
  140. {netket-3.22.3 → netket-3.23.0}/netket/driver/__init__.py +0 -0
  141. {netket-3.22.3 → netket-3.23.0}/netket/errors.py +0 -0
  142. {netket-3.22.3 → netket-3.23.0}/netket/exact.py +0 -0
  143. {netket-3.22.3 → netket-3.23.0}/netket/experimental/__init__.py +0 -0
  144. {netket-3.22.3 → netket-3.23.0}/netket/experimental/driver/__init__.py +0 -0
  145. {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/__init__.py +0 -0
  146. {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_integrator.py +0 -0
  147. {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_integrator_params.py +0 -0
  148. {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_rk/__init__.py +0 -0
  149. {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_rk/_solver.py +0 -0
  150. {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_rk/_tableau.py +0 -0
  151. {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_solver.py +0 -0
  152. {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_utils.py +0 -0
  153. {netket-3.22.3 → netket-3.23.0}/netket/experimental/geometry/__init__.py +0 -0
  154. {netket-3.22.3 → netket-3.23.0}/netket/experimental/geometry/cell.py +0 -0
  155. {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/__init__.py +0 -0
  156. {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/continuous_hilbert.py +0 -0
  157. {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/particle.py +0 -0
  158. {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/random/__init__.py +0 -0
  159. {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/random/particle.py +0 -0
  160. {netket-3.22.3 → netket-3.23.0}/netket/experimental/logging/__init__.py +0 -0
  161. {netket-3.22.3 → netket-3.23.0}/netket/experimental/models/__init__.py +0 -0
  162. {netket-3.22.3 → netket-3.23.0}/netket/experimental/models/fast_rnn.py +0 -0
  163. {netket-3.22.3 → netket-3.23.0}/netket/experimental/models/rnn.py +0 -0
  164. {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/__init__.py +0 -0
  165. {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/__init__.py +0 -0
  166. {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/cells.py +0 -0
  167. {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/layers.py +0 -0
  168. {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/layers_fast.py +0 -0
  169. {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/ordering.py +0 -0
  170. {netket-3.22.3 → netket-3.23.0}/netket/experimental/observable/__init__.py +0 -0
  171. {netket-3.22.3 → netket-3.23.0}/netket/experimental/operator/__init__.py +0 -0
  172. {netket-3.22.3 → netket-3.23.0}/netket/experimental/operator/pyscf.py +0 -0
  173. {netket-3.22.3 → netket-3.23.0}/netket/experimental/qsr/__init__.py +0 -0
  174. {netket-3.22.3 → netket-3.23.0}/netket/experimental/qsr/dataset.py +0 -0
  175. {netket-3.22.3 → netket-3.23.0}/netket/experimental/qsr/driver.py +0 -0
  176. {netket-3.22.3 → netket-3.23.0}/netket/experimental/qsr/logic_helpers.py +0 -0
  177. {netket-3.22.3 → netket-3.23.0}/netket/experimental/vqs/__init__.py +0 -0
  178. {netket-3.22.3 → netket-3.23.0}/netket/experimental/vqs/io.py +0 -0
  179. {netket-3.22.3 → netket-3.23.0}/netket/graph/__init__.py +0 -0
  180. {netket-3.22.3 → netket-3.23.0}/netket/graph/_lattice_draw.py +0 -0
  181. {netket-3.22.3 → netket-3.23.0}/netket/graph/_lattice_edge_logic.py +0 -0
  182. {netket-3.22.3 → netket-3.23.0}/netket/graph/abstract_graph.py +0 -0
  183. {netket-3.22.3 → netket-3.23.0}/netket/graph/common_lattices.py +0 -0
  184. {netket-3.22.3 → netket-3.23.0}/netket/graph/lattice.py +0 -0
  185. {netket-3.22.3 → netket-3.23.0}/netket/graph/space_group.py +0 -0
  186. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/__init__.py +0 -0
  187. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/abstract_hilbert.py +0 -0
  188. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/__init__.py +0 -0
  189. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/base.py +0 -0
  190. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/multiple.py +0 -0
  191. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/sum.py +0 -0
  192. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/sum_partition.py +0 -0
  193. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/discrete_hilbert.py +0 -0
  194. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/doubled_hilbert.py +0 -0
  195. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/fock.py +0 -0
  196. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/homogeneous.py +0 -0
  197. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/__init__.py +0 -0
  198. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/base.py +0 -0
  199. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/constrained_generic.py +0 -0
  200. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/constrained_sum.py +0 -0
  201. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/constrained_sum_partitions.py +0 -0
  202. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/unconstrained.py +0 -0
  203. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/uniform_tensor.py +0 -0
  204. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/qubit.py +0 -0
  205. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/__init__.py +0 -0
  206. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/base.py +0 -0
  207. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/doubled.py +0 -0
  208. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/fock.py +0 -0
  209. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/homogeneous.py +0 -0
  210. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/qubit.py +0 -0
  211. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/tensor_hilbert.py +0 -0
  212. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/spin.py +0 -0
  213. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/spin_orbital_fermions.py +0 -0
  214. {netket-3.22.3 → netket-3.23.0}/netket/hilbert/tensor_hilbert.py +0 -0
  215. {netket-3.22.3 → netket-3.23.0}/netket/jax/_chunk_utils.py +0 -0
  216. {netket-3.22.3 → netket-3.23.0}/netket/jax/_compat.py +0 -0
  217. {netket-3.22.3 → netket-3.23.0}/netket/jax/_expect.py +0 -0
  218. {netket-3.22.3 → netket-3.23.0}/netket/jax/_grad.py +0 -0
  219. {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/__init__.py +0 -0
  220. {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/default_mode.py +0 -0
  221. {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/jacobian_dense.py +0 -0
  222. {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/jacobian_pytree.py +0 -0
  223. {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/logic.py +0 -0
  224. {netket-3.22.3 → netket-3.23.0}/netket/jax/_math.py +0 -0
  225. {netket-3.22.3 → netket-3.23.0}/netket/jax/_ntk/__init__.py +0 -0
  226. {netket-3.22.3 → netket-3.23.0}/netket/jax/_ntk/utils.py +0 -0
  227. {netket-3.22.3 → netket-3.23.0}/netket/jax/_scanmap.py +0 -0
  228. {netket-3.22.3 → netket-3.23.0}/netket/jax/_sort.py +0 -0
  229. {netket-3.22.3 → netket-3.23.0}/netket/jax/_sparse.py +0 -0
  230. {netket-3.22.3 → netket-3.23.0}/netket/jax/_utils_dtype.py +0 -0
  231. {netket-3.22.3 → netket-3.23.0}/netket/jax/_utils_random.py +0 -0
  232. {netket-3.22.3 → netket-3.23.0}/netket/jax/_vjp.py +0 -0
  233. {netket-3.22.3 → netket-3.23.0}/netket/jax/_vjp_chunked.py +0 -0
  234. {netket-3.22.3 → netket-3.23.0}/netket/jax/_vmap_chunked.py +0 -0
  235. {netket-3.22.3 → netket-3.23.0}/netket/jax/lax.py +0 -0
  236. {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/__init__.py +0 -0
  237. {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/debug.py +0 -0
  238. {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/fixed_sharding_decorator.py +0 -0
  239. {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/fixed_sharding_utils.py +0 -0
  240. {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/flexible_sharding.py +0 -0
  241. {netket-3.22.3 → netket-3.23.0}/netket/models/README.md +0 -0
  242. {netket-3.22.3 → netket-3.23.0}/netket/models/__init__.py +0 -0
  243. {netket-3.22.3 → netket-3.23.0}/netket/models/autoreg.py +0 -0
  244. {netket-3.22.3 → netket-3.23.0}/netket/models/deepset.py +0 -0
  245. {netket-3.22.3 → netket-3.23.0}/netket/models/equivariant.py +0 -0
  246. {netket-3.22.3 → netket-3.23.0}/netket/models/fast_autoreg.py +0 -0
  247. {netket-3.22.3 → netket-3.23.0}/netket/models/full_space.py +0 -0
  248. {netket-3.22.3 → netket-3.23.0}/netket/models/gaussian.py +0 -0
  249. {netket-3.22.3 → netket-3.23.0}/netket/models/jastrow.py +0 -0
  250. {netket-3.22.3 → netket-3.23.0}/netket/models/mlp.py +0 -0
  251. {netket-3.22.3 → netket-3.23.0}/netket/models/ndm.py +0 -0
  252. {netket-3.22.3 → netket-3.23.0}/netket/models/rbm.py +0 -0
  253. {netket-3.22.3 → netket-3.23.0}/netket/models/slater.py +0 -0
  254. {netket-3.22.3 → netket-3.23.0}/netket/models/tensor_networks/__init__.py +0 -0
  255. {netket-3.22.3 → netket-3.23.0}/netket/models/tensor_networks/mpdo.py +0 -0
  256. {netket-3.22.3 → netket-3.23.0}/netket/models/tensor_networks/mps.py +0 -0
  257. {netket-3.22.3 → netket-3.23.0}/netket/models/utils.py +0 -0
  258. {netket-3.22.3 → netket-3.23.0}/netket/nn/__init__.py +0 -0
  259. {netket-3.22.3 → netket-3.23.0}/netket/nn/activation.py +0 -0
  260. {netket-3.22.3 → netket-3.23.0}/netket/nn/apply_operator.py +0 -0
  261. {netket-3.22.3 → netket-3.23.0}/netket/nn/blocks/__init__.py +0 -0
  262. {netket-3.22.3 → netket-3.23.0}/netket/nn/blocks/deepset.py +0 -0
  263. {netket-3.22.3 → netket-3.23.0}/netket/nn/blocks/mlp.py +0 -0
  264. {netket-3.22.3 → netket-3.23.0}/netket/nn/blocks/symmetry_sum.py +0 -0
  265. {netket-3.22.3 → netket-3.23.0}/netket/nn/fast_masked_linear.py +0 -0
  266. {netket-3.22.3 → netket-3.23.0}/netket/nn/masked_linear.py +0 -0
  267. {netket-3.22.3 → netket-3.23.0}/netket/nn/symmetric_linear.py +0 -0
  268. {netket-3.22.3 → netket-3.23.0}/netket/nn/utils.py +0 -0
  269. {netket-3.22.3 → netket-3.23.0}/netket/observable.py +0 -0
  270. {netket-3.22.3 → netket-3.23.0}/netket/operator/__init__.py +0 -0
  271. {netket-3.22.3 → netket-3.23.0}/netket/operator/_abstract_observable.py +0 -0
  272. {netket-3.22.3 → netket-3.23.0}/netket/operator/_abstract_operator.py +0 -0
  273. {netket-3.22.3 → netket-3.23.0}/netket/operator/_abstract_super_operator.py +0 -0
  274. {netket-3.22.3 → netket-3.23.0}/netket/operator/_bose_hubbard/__init__.py +0 -0
  275. {netket-3.22.3 → netket-3.23.0}/netket/operator/_bose_hubbard/base.py +0 -0
  276. {netket-3.22.3 → netket-3.23.0}/netket/operator/_bose_hubbard/jax.py +0 -0
  277. {netket-3.22.3 → netket-3.23.0}/netket/operator/_bose_hubbard/numba.py +0 -0
  278. {netket-3.22.3 → netket-3.23.0}/netket/operator/_continuous_operator.py +0 -0
  279. {netket-3.22.3 → netket-3.23.0}/netket/operator/_discrete_operator.py +0 -0
  280. {netket-3.22.3 → netket-3.23.0}/netket/operator/_discrete_operator_jax.py +0 -0
  281. {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/__init__.py +0 -0
  282. {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/base.py +0 -0
  283. {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/discrete_jax_operator.py +0 -0
  284. {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/discrete_operator.py +0 -0
  285. {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/operator.py +0 -0
  286. {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/__init__.py +0 -0
  287. {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/base.py +0 -0
  288. {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/jax.py +0 -0
  289. {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/numba.py +0 -0
  290. {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/utils.py +0 -0
  291. {netket-3.22.3 → netket-3.23.0}/netket/operator/_graph_operator.py +0 -0
  292. {netket-3.22.3 → netket-3.23.0}/netket/operator/_hamiltonian.py +0 -0
  293. {netket-3.22.3 → netket-3.23.0}/netket/operator/_heisenberg.py +0 -0
  294. {netket-3.22.3 → netket-3.23.0}/netket/operator/_ising/__init__.py +0 -0
  295. {netket-3.22.3 → netket-3.23.0}/netket/operator/_ising/base.py +0 -0
  296. {netket-3.22.3 → netket-3.23.0}/netket/operator/_ising/jax.py +0 -0
  297. {netket-3.22.3 → netket-3.23.0}/netket/operator/_ising/numba.py +0 -0
  298. {netket-3.22.3 → netket-3.23.0}/netket/operator/_kinetic.py +0 -0
  299. {netket-3.22.3 → netket-3.23.0}/netket/operator/_lazy.py +0 -0
  300. {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_liouvillian.py +0 -0
  301. {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/__init__.py +0 -0
  302. {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/base.py +0 -0
  303. {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/compile_helpers.py +0 -0
  304. {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/convert.py +0 -0
  305. {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/helpers.py +0 -0
  306. {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/jax.py +0 -0
  307. {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/numba.py +0 -0
  308. {netket-3.22.3 → netket-3.23.0}/netket/operator/_pauli_strings/__init__.py +0 -0
  309. {netket-3.22.3 → netket-3.23.0}/netket/operator/_pauli_strings/base.py +0 -0
  310. {netket-3.22.3 → netket-3.23.0}/netket/operator/_pauli_strings/jax.py +0 -0
  311. {netket-3.22.3 → netket-3.23.0}/netket/operator/_pauli_strings/numba.py +0 -0
  312. {netket-3.22.3 → netket-3.23.0}/netket/operator/_potential.py +0 -0
  313. {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/__init__.py +0 -0
  314. {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/base.py +0 -0
  315. {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/discrete_jax_operator.py +0 -0
  316. {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/discrete_operator.py +0 -0
  317. {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/operator.py +0 -0
  318. {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/__init__.py +0 -0
  319. {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/base.py +0 -0
  320. {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/continuous.py +0 -0
  321. {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/discrete_jax_operator.py +0 -0
  322. {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/discrete_operator.py +0 -0
  323. {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/operator.py +0 -0
  324. {netket-3.22.3 → netket-3.23.0}/netket/operator/boson.py +0 -0
  325. {netket-3.22.3 → netket-3.23.0}/netket/operator/fermion.py +0 -0
  326. {netket-3.22.3 → netket-3.23.0}/netket/operator/permutation.py +0 -0
  327. {netket-3.22.3 → netket-3.23.0}/netket/operator/spin.py +0 -0
  328. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/__init__.py +0 -0
  329. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/linear_operator.py +0 -0
  330. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/preconditioner.py +0 -0
  331. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/__init__.py +0 -0
  332. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/common.py +0 -0
  333. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/default.py +0 -0
  334. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_jacobian.py +0 -0
  335. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_jacobian_common.py +0 -0
  336. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_jacobian_dense.py +0 -0
  337. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_jacobian_pytree.py +0 -0
  338. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_onthefly.py +0 -0
  339. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_onthefly_logic.py +0 -0
  340. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/solver/__init__.py +0 -0
  341. {netket-3.22.3 → netket-3.23.0}/netket/optimizer/sr.py +0 -0
  342. {netket-3.22.3 → netket-3.23.0}/netket/sampler/__init__.py +0 -0
  343. {netket-3.22.3 → netket-3.23.0}/netket/sampler/autoreg.py +0 -0
  344. {netket-3.22.3 → netket-3.23.0}/netket/sampler/base.py +0 -0
  345. {netket-3.22.3 → netket-3.23.0}/netket/sampler/exact.py +0 -0
  346. {netket-3.22.3 → netket-3.23.0}/netket/sampler/metropolis.py +0 -0
  347. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/__init__.py +0 -0
  348. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/base.py +0 -0
  349. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/continuous_gaussian.py +0 -0
  350. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/custom_numpy.py +0 -0
  351. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/exchange.py +0 -0
  352. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/fermion_2nd.py +0 -0
  353. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/fixed.py +0 -0
  354. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/global_flip_rule.py +0 -0
  355. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/hamiltonian.py +0 -0
  356. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/hamiltonian_numpy.py +0 -0
  357. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/langevin.py +0 -0
  358. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/local.py +0 -0
  359. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/local_numpy.py +0 -0
  360. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/multiple.py +0 -0
  361. {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/tensor.py +0 -0
  362. {netket-3.22.3 → netket-3.23.0}/netket/stats/__init__.py +0 -0
  363. {netket-3.22.3 → netket-3.23.0}/netket/stats/_autocorr.py +0 -0
  364. {netket-3.22.3 → netket-3.23.0}/netket/stats/mc_stats.py +0 -0
  365. {netket-3.22.3 → netket-3.23.0}/netket/stats/mc_stats_old.py +0 -0
  366. {netket-3.22.3 → netket-3.23.0}/netket/stats/mpi_stats.py +0 -0
  367. {netket-3.22.3 → netket-3.23.0}/netket/symmetry/__init__.py +0 -0
  368. {netket-3.22.3 → netket-3.23.0}/netket/symmetry/group/__init__.py +0 -0
  369. {netket-3.22.3 → netket-3.23.0}/netket/tools/__init__.py +0 -0
  370. {netket-3.22.3 → netket-3.23.0}/netket/tools/_common.py +0 -0
  371. {netket-3.22.3 → netket-3.23.0}/netket/tools/_cpu_info.py +0 -0
  372. {netket-3.22.3 → netket-3.23.0}/netket/tools/check_mpi.py +0 -0
  373. {netket-3.22.3 → netket-3.23.0}/netket/tools/info.py +0 -0
  374. {netket-3.22.3 → netket-3.23.0}/netket/utils/__init__.py +0 -0
  375. {netket-3.22.3 → netket-3.23.0}/netket/utils/_dependencies_check.py +0 -0
  376. {netket-3.22.3 → netket-3.23.0}/netket/utils/_serialization.py +0 -0
  377. {netket-3.22.3 → netket-3.23.0}/netket/utils/api_utils.py +0 -0
  378. {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/__init__.py +0 -0
  379. {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/cli.py +0 -0
  380. {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/core.py +0 -0
  381. {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/docs.py +0 -0
  382. {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/internal_citations.bib +0 -0
  383. {netket-3.22.3 → netket-3.23.0}/netket/utils/deprecation.py +0 -0
  384. {netket-3.22.3 → netket-3.23.0}/netket/utils/dispatch.py +0 -0
  385. {netket-3.22.3 → netket-3.23.0}/netket/utils/display/__init__.py +0 -0
  386. {netket-3.22.3 → netket-3.23.0}/netket/utils/display/repr.py +0 -0
  387. {netket-3.22.3 → netket-3.23.0}/netket/utils/float.py +0 -0
  388. {netket-3.22.3 → netket-3.23.0}/netket/utils/group/__init__.py +0 -0
  389. {netket-3.22.3 → netket-3.23.0}/netket/utils/group/_group.py +0 -0
  390. {netket-3.22.3 → netket-3.23.0}/netket/utils/group/_permutation_group.py +0 -0
  391. {netket-3.22.3 → netket-3.23.0}/netket/utils/group/_point_group.py +0 -0
  392. {netket-3.22.3 → netket-3.23.0}/netket/utils/group/_semigroup.py +0 -0
  393. {netket-3.22.3 → netket-3.23.0}/netket/utils/group/axial.py +0 -0
  394. {netket-3.22.3 → netket-3.23.0}/netket/utils/group/cubic.py +0 -0
  395. {netket-3.22.3 → netket-3.23.0}/netket/utils/group/icosa.py +0 -0
  396. {netket-3.22.3 → netket-3.23.0}/netket/utils/group/planar.py +0 -0
  397. {netket-3.22.3 → netket-3.23.0}/netket/utils/history/__init__.py +0 -0
  398. {netket-3.22.3 → netket-3.23.0}/netket/utils/history/accum.py +0 -0
  399. {netket-3.22.3 → netket-3.23.0}/netket/utils/history/history_dict.py +0 -0
  400. {netket-3.22.3 → netket-3.23.0}/netket/utils/holomorphic.py +0 -0
  401. {netket-3.22.3 → netket-3.23.0}/netket/utils/iterators.py +0 -0
  402. {netket-3.22.3 → netket-3.23.0}/netket/utils/jax.py +0 -0
  403. {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/__init__.py +0 -0
  404. {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/base.py +0 -0
  405. {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/equinox.py +0 -0
  406. {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/flax.py +0 -0
  407. {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/flax_bound.py +0 -0
  408. {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/haiku.py +0 -0
  409. {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/jax.py +0 -0
  410. {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/nnx.py +0 -0
  411. {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/nnx_wrapped.py +0 -0
  412. {netket-3.22.3 → netket-3.23.0}/netket/utils/moduletools.py +0 -0
  413. {netket-3.22.3 → netket-3.23.0}/netket/utils/mpi/__init__.py +0 -0
  414. {netket-3.22.3 → netket-3.23.0}/netket/utils/mpi/mpi.py +0 -0
  415. {netket-3.22.3 → netket-3.23.0}/netket/utils/mpi/primitives.py +0 -0
  416. {netket-3.22.3 → netket-3.23.0}/netket/utils/numbers.py +0 -0
  417. {netket-3.22.3 → netket-3.23.0}/netket/utils/partial.py +0 -0
  418. {netket-3.22.3 → netket-3.23.0}/netket/utils/seed.py +0 -0
  419. {netket-3.22.3 → netket-3.23.0}/netket/utils/static_number.py +0 -0
  420. {netket-3.22.3 → netket-3.23.0}/netket/utils/static_range.py +0 -0
  421. {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/__init__.py +0 -0
  422. {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/dataclass.py +0 -0
  423. {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/pytree_serialization_sharding.py +0 -0
  424. {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/utils.py +0 -0
  425. {netket-3.22.3 → netket-3.23.0}/netket/utils/summation.py +0 -0
  426. {netket-3.22.3 → netket-3.23.0}/netket/utils/timing.py +0 -0
  427. {netket-3.22.3 → netket-3.23.0}/netket/utils/tree_walk.py +0 -0
  428. {netket-3.22.3 → netket-3.23.0}/netket/utils/types.py +0 -0
  429. {netket-3.22.3 → netket-3.23.0}/netket/utils/version_check.py +0 -0
  430. {netket-3.22.3 → netket-3.23.0}/netket/vqs/__init__.py +0 -0
  431. {netket-3.22.3 → netket-3.23.0}/netket/vqs/full_summ/__init__.py +0 -0
  432. {netket-3.22.3 → netket-3.23.0}/netket/vqs/full_summ/expect.py +0 -0
  433. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/__init__.py +0 -0
  434. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/__init__.py +0 -0
  435. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/expect.py +0 -0
  436. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/expect_chunked.py +0 -0
  437. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/expect_grad.py +0 -0
  438. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/state.py +0 -0
  439. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/__init__.py +0 -0
  440. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_forces.py +0 -0
  441. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_forces_chunked.py +0 -0
  442. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_grad.py +0 -0
  443. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_grad_chunked.py +0 -0
  444. {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/local_estimators.py +0 -0
@@ -21,6 +21,7 @@ pip-wheel-metadata
21
21
  .ipynb_checkpoints
22
22
  .pytest_cache
23
23
  .benchmarks
24
+ .jax-cache
24
25
  .mypy_cache
25
26
  .virtual_documents
26
27
  .ruff_cache
@@ -43,6 +44,8 @@ uv.lock
43
44
  # Logs and databases #
44
45
  ######################
45
46
  *.log
47
+ # per-rank output of djaxrun
48
+ djaxrunlog-p*
46
49
  *.sql
47
50
  *.sqlite
48
51
  *.json
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: NetKet
3
- Version: 3.22.3
3
+ Version: 3.23.0
4
4
  Summary: Netket : Machine Learning toolbox for many-body quantum systems.
5
5
  Project-URL: homepage, https://www.netket.org
6
6
  Project-URL: documentation, https://netket.readthedocs.io/en/latest/#
@@ -54,7 +54,7 @@ Requires-Dist: myst-nb<1.5,>=1.0; extra == 'docs'
54
54
  Requires-Dist: myst-parser<5.2.0,>=2.0.0; extra == 'docs'
55
55
  Requires-Dist: nbconvert>=7.16.6; extra == 'docs'
56
56
  Requires-Dist: sphinx-autodoc-typehints<4.0,>=1.22; extra == 'docs'
57
- Requires-Dist: sphinx-book-theme<1.3,>=1.1; extra == 'docs'
57
+ Requires-Dist: sphinx-book-theme<1.5,>=1.1; extra == 'docs'
58
58
  Requires-Dist: sphinx-copybutton>=0.5.2; extra == 'docs'
59
59
  Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
60
60
  Requires-Dist: sphinx<9.2.0,>=8.2.3; extra == 'docs'
@@ -66,6 +66,7 @@ Requires-Dist: mlflow>=2.0.0; extra == 'extra'
66
66
  Requires-Dist: openfermion>=1.0.0; extra == 'extra'
67
67
  Requires-Dist: qutip>=4; extra == 'extra'
68
68
  Requires-Dist: tensorboardx>=2.0.0; extra == 'extra'
69
+ Requires-Dist: wandb>=0.19.0; extra == 'extra'
69
70
  Provides-Extra: pyscf
70
71
  Requires-Dist: pyscf>=2.0; extra == 'pyscf'
71
72
  Description-Content-Type: text/markdown
@@ -25,6 +25,23 @@ class StopRun(Exception):
25
25
  pass
26
26
 
27
27
 
28
+ class RunFinished(StopRun):
29
+ """
30
+ Exception to be raised by callbacks to end a run that has reached its end,
31
+ for example the final time of a time evolution.
32
+
33
+ It works like :class:`~netket.callbacks.StopRun`, but the driver does not
34
+ print that the run is stopping early.
35
+ """
36
+
37
+
38
+ #: ``callback_order`` for callbacks that may raise :class:`StopRun`. They run
39
+ #: *last* (after observables at 0 and loggers at 10) so that, on the step a
40
+ #: callback stops the run, no later callback in the hook is skipped — a skipped
41
+ #: collective would deadlock the run under ``jax.distributed``.
42
+ STOPPING_CALLBACK_ORDER = 100
43
+
44
+
28
45
  class AbstractCallback(struct.Pytree, mutable=True):
29
46
  """
30
47
  Abstract base class for callbacks in advanced variational drivers.
@@ -71,7 +88,9 @@ class AbstractCallback(struct.Pytree, mutable=True):
71
88
  Lower numbers are called first, and higher numbers are called later.
72
89
 
73
90
  This can be redefined in subclasses to change the order in which callbacks are called.
74
- (Default: 0, for all callbacks, 10 for loggers).
91
+ (Default: 0, for all callbacks, 10 for loggers, and
92
+ :data:`STOPPING_CALLBACK_ORDER` (100) for callbacks that may raise
93
+ :class:`StopRun`, so they run after everything else).
75
94
  """
76
95
  return 0
77
96
 
@@ -12,13 +12,16 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from collections import deque
16
-
17
15
  import numpy as np
16
+ from flax import serialization
18
17
 
19
18
  from netket.utils import struct
20
19
 
21
- from netket._src.callbacks.base import AbstractCallback, StopRun
20
+ from netket._src.callbacks.base import (
21
+ AbstractCallback,
22
+ StopRun,
23
+ STOPPING_CALLBACK_ORDER,
24
+ )
22
25
 
23
26
 
24
27
  class ConvergenceStopping(AbstractCallback, mutable=True):
@@ -37,8 +40,11 @@ class ConvergenceStopping(AbstractCallback, mutable=True):
37
40
  """The loss must be consistently below this value for this number of
38
41
  iterations in order to stop the optimisation."""
39
42
 
40
- _loss_window: deque = struct.field(pytree_node=False, serialize=False)
41
- _patience_counter: int = struct.field(pytree_node=False, serialize=False, default=0)
43
+ # Saved, so that a run resumed from a checkpoint keeps its history. The window
44
+ # has a fixed length: only its last `_n_losses` entries are filled.
45
+ _loss_window: tuple[float, ...] = struct.field(pytree_node=False, serialize=True)
46
+ _n_losses: int = struct.field(pytree_node=False, serialize=True, default=0)
47
+ _patience_counter: int = struct.field(pytree_node=False, serialize=True, default=0)
42
48
 
43
49
  def __init__(
44
50
  self,
@@ -69,14 +75,32 @@ class ConvergenceStopping(AbstractCallback, mutable=True):
69
75
  self.smoothing_window = smoothing_window
70
76
  self.patience = patience
71
77
 
72
- self._loss_window = deque([], maxlen=smoothing_window)
78
+ self._loss_window = (0.0,) * smoothing_window
79
+ self._n_losses = 0
73
80
  self._patience_counter = 0
74
81
 
82
+ def __process_deserialization_state__(self, state):
83
+ # If the file has no window (older versions) or one of a different size,
84
+ # start the window over.
85
+ if len(state.get("_loss_window", ())) != self.smoothing_window:
86
+ state = {
87
+ k: v for k, v in state.items() if k not in ("_loss_window", "_n_losses")
88
+ }
89
+ names = ("_loss_window", "_n_losses", "_patience_counter")
90
+ defaults = {n: serialization.to_state_dict(getattr(self, n)) for n in names}
91
+ return {**defaults, **state}
92
+
93
+ @property
94
+ def callback_order(self) -> int:
95
+ # Run last, so raising StopRun never skips a later callback's collective.
96
+ return STOPPING_CALLBACK_ORDER
97
+
75
98
  def on_step_end(self, step, log_data, driver):
76
- loss = np.asarray(np.real(getattr(log_data[driver._loss_name], self.monitor)))
99
+ loss = np.real(getattr(log_data[driver._loss_name], self.monitor))
77
100
 
78
- self._loss_window.append(loss)
79
- loss_smooth = np.mean(self._loss_window)
101
+ self._loss_window = (*self._loss_window[1:], float(np.mean(loss)))
102
+ self._n_losses = min(self._n_losses + 1, self.smoothing_window)
103
+ loss_smooth = np.mean(self._loss_window[-self._n_losses :])
80
104
 
81
105
  if loss_smooth <= self.target:
82
106
  self._patience_counter += 1
@@ -16,7 +16,11 @@ import numpy as np
16
16
 
17
17
  from netket.utils import struct
18
18
 
19
- from netket._src.callbacks.base import AbstractCallback, StopRun
19
+ from netket._src.callbacks.base import (
20
+ AbstractCallback,
21
+ StopRun,
22
+ STOPPING_CALLBACK_ORDER,
23
+ )
20
24
 
21
25
 
22
26
  class EarlyStopping(AbstractCallback, mutable=True):
@@ -41,12 +45,14 @@ class EarlyStopping(AbstractCallback, mutable=True):
41
45
  start_from_step: int = struct.field(pytree_node=False)
42
46
  """Number of steps to wait before the callback has any effect."""
43
47
 
44
- _best_val: float = struct.field(pytree_node=False, serialize=False, default=np.inf)
48
+ # Saved, so that a run resumed from a checkpoint keeps counting its patience
49
+ # instead of starting over.
50
+ _best_val: float = struct.field(pytree_node=False, serialize=True, default=np.inf)
45
51
  """Best value of the loss observed up to this iteration."""
46
- _best_iter: int = struct.field(pytree_node=False, serialize=False, default=0)
52
+ _best_iter: int = struct.field(pytree_node=False, serialize=True, default=0)
47
53
  """Iteration at which the `_best_val` was observed."""
48
54
  _best_patience_counter: int = struct.field(
49
- pytree_node=False, serialize=False, default=0
55
+ pytree_node=False, serialize=True, default=0
50
56
  )
51
57
  """Stores how many steps have elapsed without an improvement."""
52
58
 
@@ -87,6 +93,19 @@ class EarlyStopping(AbstractCallback, mutable=True):
87
93
  self._best_iter = 0
88
94
  self._best_patience_counter = 0
89
95
 
96
+ # TODO: backward compatibility for checkpoints saved before September 2026.
97
+ # Remove in 2027.
98
+ def __process_deserialization_state__(self, state):
99
+ # Files saved by older versions do not contain the best-so-far state.
100
+ names = ("_best_val", "_best_iter", "_best_patience_counter")
101
+ defaults = {name: self._pytree__fields[name].default for name in names}
102
+ return {**defaults, **state}
103
+
104
+ @property
105
+ def callback_order(self) -> int:
106
+ # Run last, so raising StopRun never skips a later callback's collective.
107
+ return STOPPING_CALLBACK_ORDER
108
+
90
109
  def on_step_end(self, step, log_data, driver):
91
110
  if step < self.start_from_step:
92
111
  return
@@ -95,7 +114,7 @@ class EarlyStopping(AbstractCallback, mutable=True):
95
114
 
96
115
  self._best_patience_counter += 1
97
116
  if self._is_improvement(loss, self._best_val):
98
- self._best_val = loss
117
+ self._best_val = float(loss)
99
118
  self._best_iter = step
100
119
 
101
120
  if self.baseline is None:
@@ -16,7 +16,11 @@ import numpy as np
16
16
 
17
17
  from netket.utils import struct
18
18
 
19
- from netket._src.callbacks.base import AbstractCallback, StopRun
19
+ from netket._src.callbacks.base import (
20
+ AbstractCallback,
21
+ StopRun,
22
+ STOPPING_CALLBACK_ORDER,
23
+ )
20
24
 
21
25
 
22
26
  class InvalidLossStopping(AbstractCallback, mutable=True):
@@ -29,7 +33,9 @@ class InvalidLossStopping(AbstractCallback, mutable=True):
29
33
  patience: int | float = struct.field(pytree_node=False)
30
34
  """Number of epochs with invalid loss after which training will be stopped."""
31
35
 
32
- _last_valid_iter: int = struct.field(pytree_node=False, serialize=False, default=0)
36
+ # Saved, so that a run resumed from a checkpoint does not stop at the first
37
+ # invalid step, as if the loss had been invalid since step 0.
38
+ _last_valid_iter: int = struct.field(pytree_node=False, serialize=True, default=0)
33
39
  """Last valid iteration, to check against patience."""
34
40
 
35
41
  def __init__(self, monitor: str = "mean", patience: int | float = 0):
@@ -50,6 +56,17 @@ class InvalidLossStopping(AbstractCallback, mutable=True):
50
56
  self.patience = patience
51
57
  self._last_valid_iter = 0
52
58
 
59
+ # TODO: backward compatibility for checkpoints saved before September 2026.
60
+ # Remove in 2027.
61
+ def __process_deserialization_state__(self, state):
62
+ # Files saved by older versions do not contain it.
63
+ return {"_last_valid_iter": self._last_valid_iter, **state}
64
+
65
+ @property
66
+ def callback_order(self) -> int:
67
+ # Run last, so raising StopRun never skips a later callback's collective.
68
+ return STOPPING_CALLBACK_ORDER
69
+
53
70
  def on_step_end(self, step, log_data, driver):
54
71
  # reset if the driver was restarted
55
72
  if driver.step_count < self._last_valid_iter:
@@ -58,7 +75,9 @@ class InvalidLossStopping(AbstractCallback, mutable=True):
58
75
  if driver._loss_stats is not None:
59
76
  loss = np.real(getattr(driver._loss_stats, self.monitor))
60
77
 
61
- if not np.isfinite(loss):
78
+ # `loss` may be vector-valued (e.g. a foundation ReplicaStats has one
79
+ # mean per anchor); stop if ANY component is non-finite.
80
+ if not np.all(np.isfinite(loss)):
62
81
  if driver.step_count - self._last_valid_iter >= self.patience:
63
82
  raise StopRun(
64
83
  f"InvalidLossStopping: loss is not finite ({loss}) "
@@ -10,28 +10,39 @@ from netket.utils import struct
10
10
  from netket._src.callbacks.base import AbstractCallback
11
11
 
12
12
 
13
+ def _elapsed_time(driver, t_start: float, T_final: float) -> float:
14
+ """Time elapsed since `t_start`, at most `T_final`."""
15
+ return min(float(driver.t) - t_start, T_final)
16
+
17
+
13
18
  class TimeProgressBarCallback(AbstractCallback, mutable=True):
14
19
  """
15
20
  Progress bar displaying physical time for dynamics drivers.
16
21
 
17
- Tracks ``driver.t ∈ [0, T_final]`` instead of step count.
22
+ Tracks the time elapsed since the start of the run, ``driver.t - t_start ∈
23
+ [0, T_final]``, instead of step count. ``t_start`` is ``driver._t_start`` if
24
+ the driver sets it (as the TDVP drivers do), and 0 otherwise.
18
25
  Used automatically by :class:`~netket.driver.AbstractDynamicsDriver`
19
26
  when :meth:`run` is called with a float time argument.
20
27
  """
21
28
 
22
29
  _T_final: float = struct.field(pytree_node=False, serialize=False)
23
- _last_t: float = struct.field(pytree_node=False, serialize=False)
30
+ _t_start: float = struct.field(pytree_node=False, serialize=False)
24
31
  _pbar: Any = struct.field(pytree_node=False, serialize=False)
25
32
  _leave: bool = struct.field(pytree_node=False, serialize=False)
26
33
 
27
34
  def __init__(self, T_final: float, leave: bool = True):
28
35
  self._T_final = T_final
29
- self._last_t = 0.0
36
+ self._t_start = 0.0
30
37
  self._pbar = None
31
38
  self._leave = leave
32
39
 
33
40
  def on_run_start(self, step, driver):
34
- self._last_t = float(driver.t)
41
+ # Count from the time the run started at, even if another callback
42
+ # changes the time (e.g. when resuming from a checkpoint).
43
+ t_start = getattr(driver, "_t_start", None)
44
+ self._t_start = 0.0 if t_start is None else float(t_start)
45
+
35
46
  dt = float(driver.dt)
36
47
  n_dec = max(4, math.ceil(-math.log10(dt))) if dt > 0 else 4
37
48
  bar_format = (
@@ -46,7 +57,7 @@ class TimeProgressBarCallback(AbstractCallback, mutable=True):
46
57
  leave=self._leave,
47
58
  bar_format=bar_format,
48
59
  )
49
- self._pbar.update(float(driver.t))
60
+ self._pbar.update(_elapsed_time(driver, self._t_start, self._T_final))
50
61
  self._pbar.unpause()
51
62
  self._pbar.refresh()
52
63
 
@@ -57,16 +68,16 @@ class TimeProgressBarCallback(AbstractCallback, mutable=True):
57
68
  )
58
69
 
59
70
  def on_step_end(self, step, log_data, driver):
60
- t = float(driver.t)
61
- self._pbar.update(t - self._last_t)
62
- self._last_t = t
71
+ self._pbar.update(
72
+ _elapsed_time(driver, self._t_start, self._T_final) - self._pbar.n
73
+ )
63
74
  self._pbar.refresh()
64
75
 
65
76
  def on_run_end(self, step, driver):
66
77
  if self._pbar is not None:
67
- t = float(driver.t)
68
- self._pbar.update(t - self._last_t)
69
- self._last_t = t
78
+ self._pbar.update(
79
+ _elapsed_time(driver, self._t_start, self._T_final) - self._pbar.n
80
+ )
70
81
  self._pbar.close()
71
82
  self._pbar = None
72
83
 
@@ -80,7 +91,7 @@ class ProgressBarCallback(AbstractCallback, mutable=True):
80
91
  _leave: bool = struct.field(pytree_node=False, serialize=False)
81
92
 
82
93
  _n_steps: int = struct.field(pytree_node=False, serialize=False)
83
- _last_step: int = struct.field(pytree_node=False, serialize=False)
94
+ _start_step: int = struct.field(pytree_node=False, serialize=False)
84
95
 
85
96
  _pbar: Any = struct.field(pytree_node=False, serialize=False)
86
97
 
@@ -90,7 +101,11 @@ class ProgressBarCallback(AbstractCallback, mutable=True):
90
101
  self._leave = leave
91
102
 
92
103
  def on_run_start(self, step, driver):
93
- self._last_step = driver.step_count
104
+ # Count from the step the run started at, even if another callback
105
+ # changes the step count (e.g. when resuming from a checkpoint).
106
+ self._start_step = getattr(driver, "_start_step", None)
107
+ if self._start_step is None:
108
+ self._start_step = driver.step_count
94
109
 
95
110
  self._pbar = tqdm(
96
111
  total=self._n_steps,
@@ -98,7 +113,7 @@ class ProgressBarCallback(AbstractCallback, mutable=True):
98
113
  dynamic_ncols=True,
99
114
  leave=self._leave,
100
115
  )
101
- self._pbar.update(driver.step_count - self._last_step)
116
+ self._pbar.update(driver.step_count - self._start_step)
102
117
  self._pbar.unpause()
103
118
  self._pbar.refresh()
104
119
 
@@ -110,12 +125,12 @@ class ProgressBarCallback(AbstractCallback, mutable=True):
110
125
  )
111
126
 
112
127
  def on_step_end(self, step, log_data, driver):
113
- step_value = driver.step_count - self._last_step
128
+ step_value = driver.step_count - self._start_step
114
129
  self._pbar.update(step_value - self._pbar.n)
115
130
  self._pbar.refresh()
116
131
 
117
132
  def on_run_end(self, step, driver):
118
- step_value = driver.step_count - self._last_step
133
+ step_value = driver.step_count - self._start_step
119
134
  if self._pbar is not None:
120
135
  self._pbar.update(step_value - self._pbar.n)
121
136
  self._pbar.refresh()
@@ -1,7 +1,10 @@
1
+ import re
2
+
1
3
  from typing import Any
2
4
  from pathlib import Path
3
5
 
4
6
  import jax
7
+ import jax.experimental.multihost_utils
5
8
 
6
9
  from netket.utils import struct, timing
7
10
  from netket.utils.optional_deps import import_optional_dependency
@@ -37,7 +40,7 @@ class SaveVariationalState(AbstractCallback, mutable=True):
37
40
  >>> driver.run( # doctest: +SKIP
38
41
  ... n_iter=50,
39
42
  ... out="test",
40
- ... callback=nk.callbacks.SaveVariationalStateCallback(path="optimization", interval=10),
43
+ ... callback=nk.logging.SaveVariationalState(path="optimization", interval=10),
41
44
  ... )
42
45
  >>> nqxpack.load("optimization/state_00010.nk") # doctest: +SKIP
43
46
 
@@ -48,6 +51,9 @@ class SaveVariationalState(AbstractCallback, mutable=True):
48
51
  _file_name_root: str = struct.field(
49
52
  pytree_node=False, serialize=False, default="state"
50
53
  )
54
+ _max_to_keep: int | None = struct.field(
55
+ pytree_node=False, serialize=False, default=None
56
+ )
51
57
 
52
58
  def __init__(
53
59
  self,
@@ -55,6 +61,7 @@ class SaveVariationalState(AbstractCallback, mutable=True):
55
61
  interval: float | int,
56
62
  *,
57
63
  file_name_root: str = "state",
64
+ max_to_keep: int | None = None,
58
65
  ):
59
66
  r"""
60
67
  Constructs the callback to save the variational state at fixed intervals.
@@ -67,12 +74,15 @@ class SaveVariationalState(AbstractCallback, mutable=True):
67
74
  path (str | Path): The path where to save the variational state.
68
75
  interval (float | int, optional): The interval at which to save the variational state
69
76
  file_name_root (str, optional): The root of the file name to save the variational state. Defaults to "state".
77
+ max_to_keep (int | None, optional): If set, only the most recent ``max_to_keep`` saved states are
78
+ kept on disk; older ones are deleted after each save. Defaults to None (keep all).
70
79
  """
71
80
  if not isinstance(path, Path):
72
81
  path = Path(path)
73
82
  self._path = path
74
83
  self._interval = interval
75
84
  self._file_name_root = file_name_root
85
+ self._max_to_keep = max_to_keep
76
86
 
77
87
  if jax.process_index() == 0:
78
88
  self._path.mkdir(parents=True, exist_ok=True)
@@ -82,15 +92,35 @@ class SaveVariationalState(AbstractCallback, mutable=True):
82
92
  descr="to save the variational state with SaveVariationalStateCallback.",
83
93
  )
84
94
 
95
+ def _save(self, step, driver):
96
+ path = self._path / f"{self._file_name_root}_{step:05d}.nk"
97
+ _nqxpack_save(driver.state, path)
98
+
99
+ if self._max_to_keep is not None:
100
+ if jax.process_index() == 0:
101
+ pattern = re.compile(rf"^{re.escape(self._file_name_root)}_(\d+)\.nk$")
102
+ saved = sorted(
103
+ (
104
+ p
105
+ for p in self._path.glob(f"{self._file_name_root}_*.nk")
106
+ if pattern.match(p.name)
107
+ ),
108
+ key=lambda p: int(pattern.match(p.name).group(1)),
109
+ )
110
+ for old in saved[: -self._max_to_keep]:
111
+ old.unlink()
112
+ # Barrier so that other processes do not race ahead (e.g. saving the
113
+ # next checkpoint) while rank 0 is still pruning the directory.
114
+ jax.experimental.multihost_utils.sync_global_devices(
115
+ "SaveVariationalState._save"
116
+ )
117
+
85
118
  def before_parameter_update(self, step, log_data, driver):
86
119
  if step % self._interval == 0:
87
- path = self._path / f"{self._file_name_root}_{step:05d}.nk"
88
- _nqxpack_save(driver.state, path)
120
+ self._save(step, driver)
89
121
 
90
122
  def on_run_end(self, step, driver):
91
- if step % self._interval == 0:
92
- path = self._path / f"{self._file_name_root}_{step:05d}.nk"
93
- _nqxpack_save(driver.state, path)
123
+ self._save(step, driver)
94
124
 
95
125
 
96
126
  @timing.timed
@@ -12,7 +12,11 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from netket._src.callbacks.base import AbstractCallback, StopRun
15
+ from netket._src.callbacks.base import (
16
+ AbstractCallback,
17
+ StopRun,
18
+ STOPPING_CALLBACK_ORDER,
19
+ )
16
20
  from netket.utils import struct
17
21
 
18
22
 
@@ -30,6 +34,11 @@ class _StopAtTime(AbstractCallback):
30
34
  def __init__(self, T_final: float):
31
35
  self.T_final = T_final
32
36
 
37
+ @property
38
+ def callback_order(self) -> int:
39
+ # Run last, so raising StopRun never skips a later callback's collective.
40
+ return STOPPING_CALLBACK_ORDER
41
+
33
42
  def on_step_end(self, step, log_data, driver):
34
43
  if driver.t >= self.T_final:
35
44
  raise StopRun(f"Reached t={driver.t:.6g} >= T_final={self.T_final:.6g}")
@@ -16,7 +16,11 @@ import time
16
16
 
17
17
  from netket.utils import struct
18
18
 
19
- from netket._src.callbacks.base import AbstractCallback, StopRun
19
+ from netket._src.callbacks.base import (
20
+ AbstractCallback,
21
+ StopRun,
22
+ STOPPING_CALLBACK_ORDER,
23
+ )
20
24
 
21
25
 
22
26
  class Timeout(AbstractCallback, mutable=True):
@@ -47,11 +51,40 @@ class Timeout(AbstractCallback, mutable=True):
47
51
  self.timeout = timeout
48
52
  self._init_time = None
49
53
 
54
+ @property
55
+ def callback_order(self) -> int:
56
+ # Run last, so raising StopRun never skips a later callback's collective.
57
+ return STOPPING_CALLBACK_ORDER
58
+
50
59
  def on_run_start(self, step, driver):
51
60
  self._init_time = time.time()
52
61
 
53
62
  def on_step_end(self, step, log_data, driver):
54
- if time.time() - self._init_time >= self.timeout:
63
+ elapsed = time.time() - self._init_time
64
+ # Wall clocks and per-process start instants differ across nodes, so the
65
+ # raw per-rank decision can disagree and (without this) stop the run on
66
+ # different steps. Take process 0's decision on every rank, so the whole
67
+ # run times out in lock-step on a single authoritative clock.
68
+ if _timed_out_on_master(elapsed >= self.timeout):
55
69
  raise StopRun(
56
70
  f"Timeout: training stopped after {self.timeout:.1f} seconds."
57
71
  )
72
+
73
+
74
+ def _timed_out_on_master(local_timed_out: bool) -> bool:
75
+ """Return process 0's timeout decision on every process.
76
+
77
+ A no-op returning ``local_timed_out`` unchanged in single-process runs.
78
+ """
79
+ import jax
80
+
81
+ if jax.process_count() == 1:
82
+ return local_timed_out
83
+
84
+ import numpy as np
85
+ from jax.experimental import multihost_utils
86
+
87
+ out = multihost_utils.broadcast_one_to_all(
88
+ np.asarray(int(local_timed_out), dtype=np.int32)
89
+ )
90
+ return bool(int(out))
@@ -166,8 +166,7 @@ class AbstractDynamicsDriver(AbstractDriver):
166
166
  callbacks fire normally.
167
167
  out: Logger or iterable of loggers for output.
168
168
  obs: Observables to compute at each logging step.
169
- max_steps: Safety cap on iterations when ``T_or_n`` is a float,
170
- to prevent infinite loops if ``dt`` is zero.
169
+ n_iter: Number of iterations to perform. Mutually exclusive with ``T``.
171
170
  show_progress: Show a progress bar (default True).
172
171
  callback: User callbacks.
173
172
  timeit: If True, print timing information after the run.
@@ -31,9 +31,10 @@ class AbstractOptimizationDriver(AbstractDriver):
31
31
  Adds an `optax <https://optax.readthedocs.io/en/latest/>`_ optimizer and
32
32
  implements :meth:`update_parameters` via :func:`apply_gradient`.
33
33
 
34
- Subclass this when implementing a new optimization driver. The public
35
- name :class:`AbstractVariationalDriver` is currently kept as a compatibility
36
- alias for this class.
34
+ Subclass this when implementing a new optimization driver, following the
35
+ guidelines in :class:`~netket.driver.AbstractDriver` (including how to support
36
+ checkpointing). The public name :class:`AbstractVariationalDriver` is
37
+ currently kept as a compatibility alias for this class.
37
38
  """
38
39
 
39
40
  _optimizer: Optimizer = struct.field(pytree_node=False, serialize=False)