LICENSE.md
README.md
pyproject.toml
TACE.egg-info/PKG-INFO
TACE.egg-info/SOURCES.txt
TACE.egg-info/dependency_links.txt
TACE.egg-info/entry_points.txt
TACE.egg-info/requires.txt
TACE.egg-info/top_level.txt
eqx/LICENSE.md
eqx/README.md
eqx/__init__.py
eqx/ace/__init__.py
eqx/ace/contraction.py
eqx/ace/cuda.py
eqx/ace/tace.py
eqx/co2/__init__.py
eqx/co2/basis.py
eqx/co2/cartesian_harmonics.py
eqx/co2/gate.py
eqx/co2/irreps.py
eqx/co2/linear.py
eqx/co2/o3_tensor_product.py
eqx/co2/projection.py
eqx/co2/restriction.py
eqx/co2/spherical.py
eqx/co2/tensor_product.py
eqx/co3/__init__.py
eqx/co3/basis.py
eqx/co3/cartesian_harmonics.py
eqx/co3/gate.py
eqx/co3/irreps.py
eqx/co3/linear.py
eqx/co3/symmetric.py
eqx/co3/tensor_product.py
eqx/conv/__init__.py
eqx/conv/angular.py
eqx/conv/attention.py
eqx/conv/codegen.py
eqx/conv/contraction.py
eqx/conv/edge.py
eqx/conv/execution.py
eqx/conv/graph.py
eqx/conv/network.py
eqx/conv/program.py
eqx/conv/radial.py
eqx/conv/co3/__init__.py
eqx/conv/co3/convolution.py
eqx/conv/co3/linear.py
eqx/conv/co3/polynomials.py
eqx/conv/o2_o3/__init__.py
eqx/conv/o2_o3/autotune.py
eqx/conv/o2_o3/codegen.py
eqx/conv/o2_o3/convolution.py
eqx/conv/o2_o3/cuda.py
eqx/conv/o2_o3/direction_codegen.py
eqx/conv/o2_o3/geometry.py
eqx/conv/o2_o3/schedule.py
eqx/conv/o2_o3/transverse.py
eqx/conv/o3/__init__.py
eqx/conv/o3/codegen.py
eqx/conv/o3/convolution.py
eqx/conv/o3/cuda.py
eqx/conv/o3/harmonics.py
eqx/conv/uu_o2/__init__.py
eqx/conv/uu_o2/convolution.py
eqx/conv/uv_o2/__init__.py
eqx/conv/uv_o2/convolution.py
eqx/conv/uv_o2/transverse.py
eqx/kernels/__init__.py
eqx/kernels/channel_product.py
eqx/kernels/codegen.py
eqx/kernels/cuda.py
eqx/kernels/cuda_graph.py
eqx/kernels/layout.py
eqx/kernels/quaternion.py
eqx/kernels/recompute.py
eqx/kernels/rotary.py
eqx/kernels/rotation.py
eqx/kernels/wigner.py
eqx/kernels/csrc/runtime.cpp
eqx/models/__init__.py
eqx/models/convolution.py
eqx/models/equflash/__init__.py
eqx/models/equflash/conversion.py
eqx/models/mace/__init__.py
eqx/models/mace/conversion.py
eqx/models/nequip/__init__.py
eqx/models/nequip/conversion.py
eqx/models/prophet/__init__.py
eqx/models/prophet/conversion.py
eqx/models/sevennet/__init__.py
eqx/models/sevennet/conversion.py
eqx/models/tace/__init__.py
eqx/models/tace/tece_oam_rra/__init__.py
eqx/models/tace/tece_oam_rra/bilinear_contraction.py
eqx/models/tace/tece_oam_rra/bilinear_cuda.py
eqx/models/tace/tece_oam_rra/cuda.py
eqx/models/tace/tece_oam_rra/execution.py
eqx/models/tace/tece_oam_rra/interaction.py
eqx/models/tace/tece_oam_rra/product.py
eqx/models/tace/tece_oam_rra/program.py
eqx/nn/__init__.py
eqx/nn/activation.py
eqx/o2/__init__.py
eqx/o2/_clebsch_gordan.py
eqx/o2/_layout.py
eqx/o2/asymmetric_contraction.py
eqx/o2/circular_harmonics.py
eqx/o2/gate.py
eqx/o2/ictd.py
eqx/o2/irreps.py
eqx/o2/linear.py
eqx/o2/local_frame.py
eqx/o2/o2grid.py
eqx/o2/o3_tensor_product.py
eqx/o2/rotation_matrix.py
eqx/o2/symmetric_contraction.py
eqx/o2/tensor_product.py
eqx/o2/wigner.py
eqx/o3/__init__.py
eqx/o3/_rotation.py
eqx/o3/contraction.py
eqx/o3/cuda.py
eqx/o3/gate.py
eqx/o3/ictd.py
eqx/o3/linear.py
eqx/o3/s2grid.py
eqx/utils/__init__.py
eqx/utils/layout.py
eqx/utils/metadata.py
eqx/utils/module.py
tace/__init__.py
tace/dataset/__init__.py
tace/dataset/augmentation.py
tace/dataset/dataloader.py
tace/dataset/datamodule.py
tace/dataset/element.py
tace/dataset/graph.py
tace/dataset/neighbour_list.py
tace/dataset/quantity.py
tace/dataset/read.py
tace/dataset/sampler.py
tace/dataset/split.py
tace/dataset/statistics.py
tace/dataset/utils.py
tace/foundations/__init__.py
tace/foundations/download_link.py
tace/foundations/u_shift.py
tace/interface/__init__.py
tace/interface/ase/__init__.py
tace/interface/ase/calculator.py
tace/interface/ase/general_calculator.py
tace/interface/ase/optimizer.py
tace/interface/lammps/__init__.py
tace/interface/lammps/mliap.py
tace/interface/nvalchemi/__init__.py
tace/interface/nvalchemi/wrapper.py
tace/interface/openmm/__init__.py
tace/interface/torchsim/__init__.py
tace/interface/torchsim/torchsim.py
tace/interface/uspex26/__init__.py
tace/interface/uspex26/relax.py
tace/lightning/__init__.py
tace/lightning/lit_model.py
tace/lightning/lora.py
tace/lightning/loss_skip.py
tace/lightning/torch_model.py
tace/lightning/trainer.py
tace/lightning/u_shift.py
tace/models/__init__.py
tace/models/adapter.py
tace/models/angular.py
tace/models/blocks.py
tace/models/ictd.py
tace/models/kspace.py
tace/models/lammps.py
tace/models/layout.py
tace/models/linear.py
tace/models/mlp.py
tace/models/normalizer.py
tace/models/radial.py
tace/models/s2.py
tace/models/softmax.py
tace/models/time_reversal.py
tace/models/utils.py
tace/models/zemin.py
tace/models/_e3nn/__init__.py
tace/models/_e3nn/base.py
tace/models/_e3nn/basis_change.py
tace/models/_e3nn/default.py
tace/models/_e3nn/dropout.py
tace/models/_e3nn/edge.py
tace/models/_e3nn/fused.py
tace/models/_e3nn/inter.py
tace/models/_e3nn/layer_norm.py
tace/models/_e3nn/les.py
tace/models/_e3nn/magnetic.py
tace/models/_e3nn/node.py
tace/models/_e3nn/nonlinear.py
tace/models/_e3nn/o2.py
tace/models/_e3nn/paths.py
tace/models/_e3nn/prod.py
tace/models/_e3nn/readout.py
tace/models/_e3nn/representation.py
tace/models/_e3nn/residual.py
tace/models/_e3nn/scatter_norm.py
tace/models/_e3nn/symmetric_contraction.py
tace/models/_e3nn/tace.py
tace/models/_e3nn/tece_oam_rra.py
tace/models/_e3nn/ue.py
tace/models/_e3nn/wigner6j.py
tace/models/compile/__init__.py
tace/models/compile/aot.py
tace/models/compile/compile.py
tace/models/compile/tace.py
tace/models/compile/wrapper.py
tace/models/cue/__init__.py
tace/models/cue/_tp_scatter.py
tace/models/cue/_tp_uuu.py
tace/models/cue/paths.py
tace/models/eqt/__init__.py
tace/models/eqt/_tp_uuu.py
tace/models/eqt/equitorch/__init__.py
tace/models/eqt/equitorch/irreps/__init__.py
tace/models/eqt/equitorch/nn/__init__.py
tace/models/eqt/equitorch/nn/sparse_product.py
tace/models/eqt/equitorch/nn/tensor_products.py
tace/models/eqt/equitorch/structs/__init__.py
tace/models/oeq/__init__.py
tace/models/oeq/_tp_scatter.py
tace/scripts/__init__.py
tace/scripts/average.py
tace/scripts/clean.py
tace/scripts/convert_cgtp.py
tace/scripts/convert_lora.py
tace/scripts/eval.py
tace/scripts/export_elements.py
tace/scripts/export_eval.py
tace/scripts/export_lammps.py
tace/scripts/export_train.py
tace/scripts/finetune.py
tace/scripts/plot_diatom.py
tace/scripts/split.py
tace/scripts/train.py
tace/scripts/update.py
tace/scripts/utils/__init__.py
tace/scripts/utils/check_equi.py
tace/scripts/utils/check_soc.py
tace/scripts/utils/convert_to_xyz.py
tace/utils/__init__.py
tace/utils/_global.py
tace/utils/callbacks.py
tace/utils/cfg.py
tace/utils/ema.py
tace/utils/env.py
tace/utils/hydra_resolver.py
tace/utils/logger.py
tace/utils/mask_metrics.py
tace/utils/metrics.py
tace/utils/metrics_bake.py
tace/utils/spectra.py
tace/utils/strategy.py
tace/utils/torch_scatter.py
tace/utils/units.py
tace/utils/utils.py
tace/utils/loss/__init__.py
tace/utils/loss/common.py
tace/utils/loss/dens.py
tace/utils/loss/huber_fn.py
tace/utils/loss/l2mae_fn.py
tace/utils/loss/mae_fn.py
tace/utils/loss/mse_fn.py
tace/utils/loss/normal.py
tace/utils/loss/registry.py
tace/utils/loss/special_fn.py
tace/utils/loss/uncertainty.py
tace/utils/lr_scheduler/__init__.py
tace/utils/lr_scheduler/warmup.py
tace/utils/lr_scheduler/wsd.py
tace/utils/optimizer/__init__.py
tace/utils/optimizer/hybrid_muon.py
tace/utils/optimizer/soap.py
tests/test_cartesian.py
tests/test_compile.py
tests/test_eqx_integration.py
tests/test_export_elements.py
tests/test_interfaces.py
tests/test_les.py
tests/test_local_o2.py
tests/test_loss.py
tests/test_magnetic_optimizer.py
tests/test_product.py
tests/test_radial.py
tests/test_scatter_norm.py
tests/test_statistics.py
tests/test_time_reversal.py
tests/test_torchsim.py
tests/test_wigner6j.py