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.
- {netket-3.22.3 → netket-3.23.0}/.gitignore +3 -0
- {netket-3.22.3 → netket-3.23.0}/PKG-INFO +4 -3
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/base.py +20 -1
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/convergence_stopping.py +33 -9
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/early_stopping.py +24 -5
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/invalid_loss_stopping.py +22 -3
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/progressbar.py +31 -16
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/save_state.py +36 -6
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/stop_at_time.py +10 -1
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/timeout.py +35 -2
- {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/abstract_dynamics_driver.py +1 -2
- {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/abstract_optimization_driver.py +4 -3
- {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/abstract_variational_driver.py +36 -12
- {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/infidelity_sr.py +7 -2
- {netket-3.22.3 → netket-3.23.0}/netket/_src/driver/vmc_sr.py +56 -9
- {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/hdf5_log.py +17 -12
- {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/mlflow_log.py +26 -5
- {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/tensorboard.py +15 -3
- netket-3.23.0/netket/_src/logging/wandb_log.py +376 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/srt_onthefly.py +53 -8
- {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/common.py +25 -5
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/pyscf_api.py +1 -1
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/pyscf_utils.py +1 -1
- netket-3.23.0/netket/_src/solvers/jaxmg_interface.py +226 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/solvers/nan_fallback.py +8 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/operations.py +1 -1
- {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/representation.py +1 -1
- netket-3.23.0/netket/_src/vqs/expect_to_precision.py +221 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/freeze_vstate.py +16 -50
- {netket-3.22.3 → netket-3.23.0}/netket/_version.py +2 -2
- {netket-3.22.3 → netket-3.23.0}/netket/callbacks.py +1 -0
- {netket-3.22.3 → netket-3.23.0}/netket/driver/steady_state.py +4 -6
- {netket-3.22.3 → netket-3.23.0}/netket/driver/vmc.py +4 -6
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/driver/tdvp.py +13 -0
- netket-3.23.0/netket/experimental/driver/tdvp_common.py +518 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/driver/tdvp_schmitt.py +17 -1
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_integrator_state.py +5 -2
- {netket-3.22.3 → netket-3.23.0}/netket/graph/graph.py +1 -1
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/tensor_hilbert_discrete.py +1 -1
- {netket-3.22.3 → netket-3.23.0}/netket/jax/__init__.py +1 -0
- netket-3.23.0/netket/jax/_apply_packed.py +461 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_ntk/logic.py +18 -3
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_utils_tree.py +16 -0
- {netket-3.22.3 → netket-3.23.0}/netket/logging.py +1 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/solver/solvers.py +116 -64
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/metropolis_numpy.py +13 -2
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/parallel_tempering.py +8 -2
- {netket-3.22.3 → netket-3.23.0}/netket/tools/djaxrun.py +28 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/_tips.py +1 -1
- {netket-3.22.3 → netket-3.23.0}/netket/utils/array.py +1 -1
- {netket-3.22.3 → netket-3.23.0}/netket/utils/config_flags.py +24 -4
- {netket-3.22.3 → netket-3.23.0}/netket/utils/history/history.py +24 -1
- netket-3.23.0/netket/utils/optional_deps.py +81 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/fields.py +5 -2
- {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/pytree.py +24 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/base.py +26 -4
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/full_summ/state.py +11 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/common.py +13 -4
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/kernels.py +1 -1
- netket-3.23.0/netket/vqs/mc/kernels_packed.py +212 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect.py +4 -1
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_chunked.py +12 -3
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/state.py +62 -14
- {netket-3.22.3 → netket-3.23.0}/pyproject.toml +3 -3
- netket-3.22.3/netket/_src/driver/abstract_variational_driver_old.py +0 -443
- netket-3.22.3/netket/_src/vqs/expect_to_precision.py +0 -164
- netket-3.22.3/netket/experimental/driver/tdvp_common.py +0 -457
- netket-3.22.3/netket/utils/optional_deps.py +0 -39
- {netket-3.22.3 → netket-3.23.0}/LICENSE +0 -0
- {netket-3.22.3 → netket-3.23.0}/README.md +0 -0
- {netket-3.22.3 → netket-3.23.0}/docs/sphinx_extensions/netket_citations.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/auto_chunk_size.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/auto_slurm_requeue.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/callback_list.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/legacy_wrappers.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/callbacks/observables.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/json_log.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/runtime_log.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/logging/state_log.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/kwargs.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/sr.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/sr_srt_common.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/ngd/srt.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/apply_operator/functional.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/apply_operator/linen.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/apply_operator/nnx.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/api.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/functional.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/linen.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/nn/freeze/nnx.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/correlator/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/correlator/correlator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/correlator/expect.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/infidelity/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/infidelity/exact.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/infidelity/expect.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/infidelity/infidelity_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/renyi2/S2_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/renyi2/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/renyi2/exact.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/renyi2/expect.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/variance/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/variance/exact.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/variance/expect.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/variance/variance_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/observable/vscore.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/normal_order_utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/conversion.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/fermihubbard.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/kernels.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/operator_data.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/particle_number_conserving_fermionic/operators.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/construct.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/fermion.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/spin.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/permutation/trace_utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/operator/spin_flip.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/local_estimators.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/accumulator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/accumulator_batch.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/stats/online_stats/kernels.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/canonical_representation.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/labeled_representation.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/labeled_representation_coset_filter.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/representation_construction.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/spin_flip_representation.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/translation_coset_filter.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/symmetry/translation_representation.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/check_mc_convergence.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/fermion_mf/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/fermion_mf/expectation.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/fermion_mf/qgt.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/fermion_mf/state.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/plot_mc_convergence.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/_src/vqs/transformed_vstate.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/driver/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/errors.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/exact.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/driver/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_integrator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_integrator_params.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_rk/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_rk/_solver.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_rk/_tableau.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_solver.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/dynamics/_utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/geometry/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/geometry/cell.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/continuous_hilbert.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/particle.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/random/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/hilbert/random/particle.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/logging/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/models/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/models/fast_rnn.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/models/rnn.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/cells.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/layers.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/layers_fast.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/nn/rnn/ordering.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/observable/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/operator/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/operator/pyscf.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/qsr/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/qsr/dataset.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/qsr/driver.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/qsr/logic_helpers.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/vqs/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/experimental/vqs/io.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/graph/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/graph/_lattice_draw.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/graph/_lattice_edge_logic.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/graph/abstract_graph.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/graph/common_lattices.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/graph/lattice.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/graph/space_group.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/abstract_hilbert.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/multiple.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/sum.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/constraint/sum_partition.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/discrete_hilbert.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/doubled_hilbert.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/fock.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/homogeneous.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/constrained_generic.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/constrained_sum.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/constrained_sum_partitions.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/unconstrained.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/index/uniform_tensor.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/qubit.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/doubled.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/fock.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/homogeneous.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/qubit.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/random/tensor_hilbert.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/spin.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/spin_orbital_fermions.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/hilbert/tensor_hilbert.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_chunk_utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_compat.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_expect.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_grad.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/default_mode.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/jacobian_dense.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/jacobian_pytree.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_jacobian/logic.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_math.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_ntk/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_ntk/utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_scanmap.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_sort.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_sparse.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_utils_dtype.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_utils_random.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_vjp.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_vjp_chunked.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/_vmap_chunked.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/lax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/debug.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/fixed_sharding_decorator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/fixed_sharding_utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/jax/sharding/flexible_sharding.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/README.md +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/autoreg.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/deepset.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/equivariant.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/fast_autoreg.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/full_space.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/gaussian.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/jastrow.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/mlp.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/ndm.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/rbm.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/slater.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/tensor_networks/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/tensor_networks/mpdo.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/tensor_networks/mps.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/models/utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/activation.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/apply_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/blocks/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/blocks/deepset.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/blocks/mlp.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/blocks/symmetry_sum.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/fast_masked_linear.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/masked_linear.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/symmetric_linear.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/nn/utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/observable.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_abstract_observable.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_abstract_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_abstract_super_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_bose_hubbard/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_bose_hubbard/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_bose_hubbard/jax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_bose_hubbard/numba.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_continuous_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_discrete_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_discrete_operator_jax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/discrete_jax_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/discrete_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_embed/operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/jax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/numba.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_fermion2nd/utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_graph_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_hamiltonian.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_heisenberg.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_ising/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_ising/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_ising/jax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_ising/numba.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_kinetic.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_lazy.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_liouvillian.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/compile_helpers.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/convert.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/helpers.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/jax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_local_operator/numba.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_pauli_strings/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_pauli_strings/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_pauli_strings/jax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_pauli_strings/numba.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_potential.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/discrete_jax_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/discrete_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_prod/operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/continuous.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/discrete_jax_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/discrete_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/_sum/operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/boson.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/fermion.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/permutation.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/operator/spin.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/linear_operator.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/preconditioner.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/common.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/default.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_jacobian.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_jacobian_common.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_jacobian_dense.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_jacobian_pytree.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_onthefly.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/qgt/qgt_onthefly_logic.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/solver/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/optimizer/sr.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/autoreg.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/exact.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/metropolis.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/continuous_gaussian.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/custom_numpy.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/exchange.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/fermion_2nd.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/fixed.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/global_flip_rule.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/hamiltonian.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/hamiltonian_numpy.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/langevin.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/local.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/local_numpy.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/multiple.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/sampler/rules/tensor.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/stats/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/stats/_autocorr.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/stats/mc_stats.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/stats/mc_stats_old.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/stats/mpi_stats.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/symmetry/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/symmetry/group/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/tools/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/tools/_common.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/tools/_cpu_info.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/tools/check_mpi.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/tools/info.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/_dependencies_check.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/_serialization.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/api_utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/cli.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/core.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/docs.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/citations/internal_citations.bib +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/deprecation.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/dispatch.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/display/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/display/repr.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/float.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/group/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/group/_group.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/group/_permutation_group.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/group/_point_group.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/group/_semigroup.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/group/axial.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/group/cubic.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/group/icosa.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/group/planar.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/history/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/history/accum.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/history/history_dict.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/holomorphic.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/iterators.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/jax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/base.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/equinox.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/flax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/flax_bound.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/haiku.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/jax.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/nnx.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/model_frameworks/nnx_wrapped.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/moduletools.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/mpi/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/mpi/mpi.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/mpi/primitives.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/numbers.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/partial.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/seed.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/static_number.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/static_range.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/dataclass.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/pytree_serialization_sharding.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/struct/utils.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/summation.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/timing.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/tree_walk.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/types.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/utils/version_check.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/full_summ/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/full_summ/expect.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/expect.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/expect_chunked.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/expect_grad.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_mixed_state/state.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/__init__.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_forces.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_forces_chunked.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_grad.py +0 -0
- {netket-3.22.3 → netket-3.23.0}/netket/vqs/mc/mc_state/expect_grad_chunked.py +0 -0
- {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.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: NetKet
|
|
3
|
-
Version: 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.
|
|
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
|
|
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
|
-
|
|
41
|
-
|
|
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 =
|
|
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.
|
|
99
|
+
loss = np.real(getattr(log_data[driver._loss_name], self.monitor))
|
|
77
100
|
|
|
78
|
-
self._loss_window.
|
|
79
|
-
|
|
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
|
|
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
|
-
|
|
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=
|
|
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=
|
|
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
|
|
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
|
-
|
|
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
|
-
|
|
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
|
|
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
|
-
|
|
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.
|
|
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
|
-
|
|
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(
|
|
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
|
-
|
|
61
|
-
|
|
62
|
-
|
|
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
|
-
|
|
68
|
-
|
|
69
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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.
|
|
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.
|
|
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.
|
|
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.
|
|
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
|
-
|
|
88
|
-
_nqxpack_save(driver.state, path)
|
|
120
|
+
self._save(step, driver)
|
|
89
121
|
|
|
90
122
|
def on_run_end(self, step, driver):
|
|
91
|
-
|
|
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
|
|
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
|
|
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
|
-
|
|
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
|
-
|
|
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
|
|
35
|
-
|
|
36
|
-
|
|
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)
|