LICENSE
NOTICE
README.md
THIRD_PARTY_NOTICES.md
pyproject.toml
src/hamiltonzero/__init__.py
src/hamiltonzero/checkpoint.py
src/hamiltonzero/cli.py
src/hamiltonzero/config.py
src/hamiltonzero/hamiltonian.py
src/hamiltonzero/inference.py
src/hamiltonzero/observables.py
src/hamiltonzero/renyi2.py
src/hamiltonzero.egg-info/PKG-INFO
src/hamiltonzero.egg-info/SOURCES.txt
src/hamiltonzero.egg-info/dependency_links.txt
src/hamiltonzero.egg-info/entry_points.txt
src/hamiltonzero.egg-info/requires.txt
src/hamiltonzero.egg-info/top_level.txt
src/hamiltonzero/compiled/__init__.py
src/hamiltonzero/compiled/api.py
src/hamiltonzero/compiled/execute.py
src/hamiltonzero/compiled/model.py
src/hamiltonzero/compiled/tree.py
src/hamiltonzero/compiled/trunk.py
src/hamiltonzero/compiled/types.py
src/hamiltonzero/data/__init__.py
src/hamiltonzero/data/systems.py
src/hamiltonzero/energy/__init__.py
src/hamiltonzero/energy/compiled.py
src/hamiltonzero/energy/custom_lap.py
src/hamiltonzero/energy/frame.py
src/hamiltonzero/energy/kernel.py
src/hamiltonzero/evaluation/__init__.py
src/hamiltonzero/evaluation/backend.py
src/hamiltonzero/evaluation/greedy_router.py
src/hamiltonzero/evaluation/large_n.py
src/hamiltonzero/evaluation/pallas_mha.py
src/hamiltonzero/evaluation/runner.py
src/hamiltonzero/evaluation/runtime.py
src/hamiltonzero/evaluation/sequence_parallel.py
src/hamiltonzero/evaluation/sequence_trunk.py
src/hamiltonzero/evaluation/statistics.py
src/hamiltonzero/evaluation/types.py
src/hamiltonzero/mcmc/__init__.py
src/hamiltonzero/mcmc/_monotone_cubic.py
src/hamiltonzero/mcmc/langevin.py
src/hamiltonzero/mcmc/quaternion.py
src/hamiltonzero/mcmc/replica_exchange.py
src/hamiltonzero/mcmc/runtime.py
src/hamiltonzero/model/__init__.py
src/hamiltonzero/model/_custom_lap_primitives.py
src/hamiltonzero/model/api.py
src/hamiltonzero/model/context.py
src/hamiltonzero/model/featurizer.py
src/hamiltonzero/model/fp32.py
src/hamiltonzero/model/fused_silu.py
src/hamiltonzero/model/global_ladder.py
src/hamiltonzero/model/model.py
src/hamiltonzero/model/odd_ops.py
src/hamiltonzero/model/pallas_attention.py
src/hamiltonzero/model/readout_leaf_context.py
src/hamiltonzero/model/route_pointer.py
src/hamiltonzero/model/route_quotient.py
src/hamiltonzero/model/tree.py
src/hamiltonzero/model/trunk.py
src/hamiltonzero/model/_pallas_attn/__init__.py
src/hamiltonzero/model/_pallas_attn/custom_gradients.py
src/hamiltonzero/model/_pallas_attn/mhsea.py
src/hamiltonzero/model/_pallas_attn/mhsea_full_ad.py
src/hamiltonzero/model/_pallas_attn/mhsea_full_ad_jvp.py
src/hamiltonzero/model/_pallas_attn/mhsea_tuned.py
src/hamiltonzero/model/_pallas_attn/mhsea_tuned_jvp.py
src/hamiltonzero/model/_pallas_attn/utils.py
src/hamiltonzero/modes/__init__.py
src/hamiltonzero/modes/eval.py
src/hamiltonzero/modes/finetune.py
src/hamiltonzero/modes/train.py
src/hamiltonzero/optim/__init__.py
src/hamiltonzero/optim/api.py
src/hamiltonzero/optim/blocks.py
src/hamiltonzero/optim/compat.py
src/hamiltonzero/optim/production.py
src/hamiltonzero/optim/spin_blocks.py
src/hamiltonzero/optim/targets.py
src/hamiltonzero/router/__init__.py
src/hamiltonzero/router/api.py
src/hamiltonzero/router/baseline.py
src/hamiltonzero/router/compiled.py
src/hamiltonzero/router/decode.py
src/hamiltonzero/router/permutation.py
src/hamiltonzero/router/state.py
src/hamiltonzero/router/types.py
src/kfac_jax/__init__.py
src/kfac_jax/py.typed
src/kfac_jax/_src/__init__.py
src/kfac_jax/_src/layers_and_loss_tags.py
src/kfac_jax/_src/loss_functions.py
src/kfac_jax/_src/optimizer.py
src/kfac_jax/_src/patches_second_moment.py
src/kfac_jax/_src/tag_graph_matcher.py
src/kfac_jax/_src/tracer.py
src/kfac_jax/_src/curvature_blocks/__init__.py
src/kfac_jax/_src/curvature_blocks/curvature_block.py
src/kfac_jax/_src/curvature_blocks/diagonal.py
src/kfac_jax/_src/curvature_blocks/full.py
src/kfac_jax/_src/curvature_blocks/kronecker_factored.py
src/kfac_jax/_src/curvature_blocks/tnt.py
src/kfac_jax/_src/curvature_blocks/utils.py
src/kfac_jax/_src/curvature_estimator/__init__.py
src/kfac_jax/_src/curvature_estimator/block_diagonal.py
src/kfac_jax/_src/curvature_estimator/curvature_estimator.py
src/kfac_jax/_src/curvature_estimator/explicit_exact.py
src/kfac_jax/_src/curvature_estimator/implicit_exact.py
src/kfac_jax/_src/curvature_estimator/optax_interface.py
src/kfac_jax/_src/utils/__init__.py
src/kfac_jax/_src/utils/accumulators.py
src/kfac_jax/_src/utils/math.py
src/kfac_jax/_src/utils/misc.py
src/kfac_jax/_src/utils/parallel.py
src/kfac_jax/_src/utils/staging.py
src/kfac_jax/_src/utils/types.py
third_party/folx/LICENSE
third_party/jax/LICENSE
third_party/kfac_jax/LICENSE