Index _ | A | B | C | D | E | F | G | H | I | J | K | L | M | N | O | P | R | S | T | U | V | W | Z _ __init__() (nvalchemi.dynamics.hooks.ConvergedSnapshotHook method) (nvalchemi.dynamics.hooks.EnergyDriftMonitorHook method) (nvalchemi.dynamics.hooks.FreezeAtomsHook method) (nvalchemi.dynamics.hooks.LoggingHook method) (nvalchemi.dynamics.hooks.MaxForceClampHook method) (nvalchemi.dynamics.hooks.NaNDetectorHook method) (nvalchemi.dynamics.hooks.SnapshotHook method) (nvalchemi.dynamics.hooks.StageTimingHook method), [1] (nvalchemi.dynamics.hooks.TorchProfilerHook method), [1] (nvalchemi.training.hooks.TorchProfilerHook method) __len__() (nvalchemi.dynamics.DataSink method) __needs_keys__ (nvalchemi.dynamics.BaseDynamics attribute) (nvalchemi.dynamics.DemoDynamics attribute) (nvalchemi.dynamics.FusedStage attribute) __provides_keys__ (nvalchemi.dynamics.BaseDynamics attribute) (nvalchemi.dynamics.DemoDynamics attribute) (nvalchemi.dynamics.FusedStage attribute) _config (nvalchemi.data.AtomicDataZarrWriter attribute) _dist_initialized (nvalchemi.dynamics.DistributedPipeline attribute) _exclusive_update_key (nvalchemi.training.hooks.TrainingUpdateHook attribute) _prev_accelerations (nvalchemi.dynamics.DemoDynamics attribute) _profiled_stages (nvalchemi.dynamics.hooks.StageTimingHook attribute), [1] (nvalchemi.hooks.StageTimingHook attribute) _stage_type (nvalchemi.hooks.HookRegistryMixin attribute) _store (nvalchemi.data.AtomicDataZarrReader attribute) (nvalchemi.data.AtomicDataZarrWriter attribute) A action (nvalchemi.dynamics.hooks.EnergyDriftMonitorHook attribute), [1] active_outputs (nvalchemi.models.base.ModelConfig attribute) activities (nvalchemi.hooks.TorchProfilerHook attribute) adapt_input() (nvalchemi.models.aimnet2.AIMNet2Wrapper method) (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.demo.DemoModelWrapper method) (nvalchemi.models.dftd3.DFTD3ModelWrapper method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.lj.LennardJonesModelWrapper method) (nvalchemi.models.mace.MACEWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) (nvalchemi.models.uma.UMAWrapper method) adapt_output() (nvalchemi.models.aimnet2.AIMNet2Wrapper method) (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.demo.DemoModelWrapper method) (nvalchemi.models.dftd3.DFTD3ModelWrapper method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.lj.LennardJonesModelWrapper method) (nvalchemi.models.mace.MACEWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) (nvalchemi.models.uma.UMAWrapper method) AdapterRegistry (class in nvalchemi.distributed) AdapterStatus (class in nvalchemi.distributed) add_custom() (nvalchemi.data.AtomicDataZarrWriter method) add_key() (nvalchemi.data.Batch method) add_message() (nvalchemi.hooks.ReportingState method) add_output_head() (nvalchemi.models.base.BaseModelMixin method) AFTER_COMPUTE (nvalchemi.dynamics.DynamicsStage attribute) AFTER_POST_UPDATE (nvalchemi.dynamics.DynamicsStage attribute) AFTER_PRE_UPDATE (nvalchemi.dynamics.DynamicsStage attribute) AFTER_STEP (nvalchemi.dynamics.DynamicsStage attribute) AIMNet2Wrapper (class in nvalchemi.models.aimnet2) align_dtypes() (nvalchemi.training.BaseLossFunction method) all_complete() (nvalchemi.dynamics.FusedStage static method) all_reduce_replicated_outputs (nvalchemi.distributed.OpAdapter property) append() (nvalchemi.data.AtomicDataZarrWriter method) (nvalchemi.data.Batch method) append_data() (nvalchemi.data.Batch method) as_dict() (nvalchemi.hooks.ScalarSnapshot method) assign_atoms_to_ranks() (nvalchemi.distributed.SpatialPartitioner method) associations (nvalchemi.training.CheckpointManifest attribute) async_save (nvalchemi.training.hooks.CheckpointHook attribute) atom_categories (nvalchemi.data.AtomicData attribute) atom_fields() (nvalchemi.distributed.ShardedBatch method) atomic_masses (nvalchemi.data.AtomicData attribute) atomic_numbers (nvalchemi.data.AtomicData attribute) AtomicDataZarrReader (class in nvalchemi.data) AtomicDataZarrWriter (class in nvalchemi.data) auto_init (nvalchemi.training.hooks.DDPHook attribute) autograd_inputs (nvalchemi.models.base.ModelConfig attribute) autograd_outputs (nvalchemi.models.base.ModelConfig attribute) autograd_target() (in module nvalchemi.distributed) B backend (nvalchemi.training.hooks.DDPHook attribute) balance_cells_for_ranks() (nvalchemi.distributed.SpatialPartitioner static method) balanced() (nvalchemi.data.datapipes.MultiDatasetBatchSampler class method) BaseDynamics (class in nvalchemi.dynamics) BaseLossFunction (class in nvalchemi.training) BaseModelMixin (class in nvalchemi.models.base) BaseRichLayout (class in nvalchemi.hooks) Batch (class in nvalchemi.data) batch (nvalchemi.hooks.HookContext attribute) batch_callback (nvalchemi.training.ValidationConfig attribute) batch_count (nvalchemi.hooks.ScalarSnapshot attribute) (nvalchemi.hooks.TrainContext attribute) (nvalchemi.training.TrainingStrategy attribute) batch_idx (nvalchemi.data.Batch property) batch_ptr (nvalchemi.data.Batch property) batch_size (nvalchemi.data.Batch property) (nvalchemi.data.DataLoader attribute) BatchValidationCallback (class in nvalchemi.training) BEFORE_COMPUTE (nvalchemi.dynamics.DynamicsStage attribute) BEFORE_POST_UPDATE (nvalchemi.dynamics.DynamicsStage attribute) BEFORE_PRE_UPDATE (nvalchemi.dynamics.DynamicsStage attribute) BEFORE_STEP (nvalchemi.dynamics.DynamicsStage attribute) bias_fn (nvalchemi.hooks.BiasedPotentialHook attribute) BiasedPotentialHook (class in nvalchemi.hooks) boundaries (nvalchemi.training.PiecewiseWeight attribute) broadcast_buffers (nvalchemi.distributed.DistributedManager property) (nvalchemi.training.hooks.DDPHook attribute) build_initial_batch() (nvalchemi.dynamics.SizeAwareSampler method) C cancel_prefetch() (nvalchemi.data.datapipes.MultiDataset method) (nvalchemi.data.Dataset method) (nvalchemi.data.InMemoryDataset method) capacity (nvalchemi.dynamics.DataSink attribute) (nvalchemi.dynamics.DataSink property) (nvalchemi.dynamics.GPUBuffer attribute) (nvalchemi.dynamics.GPUBuffer property) (nvalchemi.dynamics.HostMemory attribute) (nvalchemi.dynamics.HostMemory property) (nvalchemi.dynamics.ZarrData attribute) (nvalchemi.dynamics.ZarrData property) cell (nvalchemi.data.AtomicData attribute) cell_to_rank() (nvalchemi.distributed.SpatialPartitioner method) charge (nvalchemi.data.AtomicData attribute) charges (nvalchemi.data.AtomicData attribute) checkpoint_dir (nvalchemi.training.hooks.CheckpointHook attribute) checkpoint_index (nvalchemi.training.CheckpointManifest attribute) checkpoint_spec() (nvalchemi.models.mace.MACEWrapper method) CheckpointableHook (class in nvalchemi.hooks) chunk_size (nvalchemi.data.datapipes.ZarrArrayConfig attribute) cleanup() (nvalchemi.distributed.DistributedManager static method) (nvalchemi.dynamics.DistributedPipeline method) clone() (nvalchemi.data.Batch method) close() (nvalchemi.data.AtomicDataZarrReader method) (nvalchemi.data.datapipes.MultiDataset method) (nvalchemi.data.Dataset method) (nvalchemi.data.InMemoryDataset method) (nvalchemi.data.Reader method) (nvalchemi.distributed.DistributedModel method) (nvalchemi.distributed.DistributedPipelineModel method) (nvalchemi.distributed.DomainParallel method) (nvalchemi.dynamics.hooks.LoggingHook method) (nvalchemi.hooks.ReportingOrchestrator method) (nvalchemi.hooks.RichReporter method) (nvalchemi.hooks.StageTimingHook method) (nvalchemi.hooks.TensorBoardReporter method) (nvalchemi.training.hooks.TrainingUpdateOrchestrator method) cls_path (nvalchemi.training.BaseSpec attribute) collect_scalars() (in module nvalchemi.hooks) collective_device() (in module nvalchemi.distributed) compile() (nvalchemi.dynamics.FusedStage method) compile_kwargs (nvalchemi.dynamics.FusedStage attribute) compile_step (nvalchemi.dynamics.FusedStage attribute) CompilePolicy (class in nvalchemi.distributed) components (nvalchemi.training.ComposedLossFunction attribute) ComposedLossFunction (class in nvalchemi.training) ComposedLossOutput (class in nvalchemi.training) compressors (nvalchemi.data.datapipes.ZarrArrayConfig attribute) compute() (nvalchemi.dynamics.BaseDynamics method) compute_barostat_mass() (in module nvalchemiops.dynamics.integrators) compute_barostat_potential_energy() (in module nvalchemiops.dynamics.integrators) compute_cell_kinetic_energy() (in module nvalchemiops.dynamics.integrators) compute_embeddings() (nvalchemi.models.aimnet2.AIMNet2Wrapper method) (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.demo.DemoModelWrapper method) (nvalchemi.models.dftd3.DFTD3ModelWrapper method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.lj.LennardJonesModelWrapper method) (nvalchemi.models.mace.MACEWrapper method) (nvalchemi.models.pipeline.PipelineModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) (nvalchemi.models.uma.UMAWrapper method) compute_kinetic_tensor() (in module nvalchemiops.dynamics.integrators) compute_pressure_tensor() (in module nvalchemiops.dynamics.integrators) compute_rank_grid() (nvalchemi.distributed.SpatialPartitioner static method) compute_residual() (nvalchemi.training.BaseLossFunction method) (nvalchemi.training.EnergyHuberLoss method) (nvalchemi.training.EnergyMAELoss method) (nvalchemi.training.EnergyMSELoss method) (nvalchemi.training.ForceHuberLoss method) (nvalchemi.training.ForceL2NormLoss method) (nvalchemi.training.ForceMSELoss method) (nvalchemi.training.StressHuberLoss method) (nvalchemi.training.StressMSELoss method) compute_scalar_pressure() (in module nvalchemiops.dynamics.integrators) config (nvalchemi.distributed.DistributedModel property) consolidate_node_energy (nvalchemi.distributed.CompilePolicy property) contiguous() (nvalchemi.data.Batch method) converged_mask (nvalchemi.hooks.DynamicsContext attribute) ConvergedSnapshotHook (class in nvalchemi.dynamics.hooks), [1] convergence_hook (nvalchemi.dynamics.BaseDynamics attribute) ConvergenceHook (class in nvalchemi.dynamics) COOPadder (class in nvalchemi.distributed) core (nvalchemi.data.datapipes.ZarrWriteConfig attribute) cpu() (nvalchemi.data.Batch method) create_model_spec() (in module nvalchemi.training) create_model_spec_from_json() (in module nvalchemi.training) create_orthogonal_process_group() (nvalchemi.distributed.DistributedManager static method) create_process_subgroup() (nvalchemi.distributed.DistributedManager static method) criteria (nvalchemi.dynamics.ConvergenceHook attribute) cuda (nvalchemi.distributed.DistributedManager property) cuda() (nvalchemi.data.Batch method) current_dd_context() (in module nvalchemi.distributed) current_weight() (nvalchemi.training.ComposedLossFunction method) custom (nvalchemi.data.datapipes.ZarrWriteConfig attribute) cutoff (nvalchemi.models.base.NeighborConfig attribute) (nvalchemi.models.mace.MACEWrapper property) (nvalchemi.models.uma.UMAWrapper property) D DataLoader (class in nvalchemi.data) Dataset (class in nvalchemi.data) dataset (nvalchemi.data.DataLoader attribute) datasets (nvalchemi.data.datapipes.MultiDataset property) DataSink (class in nvalchemi.dynamics) default_preview_batch_count() (nvalchemi.hooks.BaseRichLayout method) (nvalchemi.hooks.DynamicsRichLayout method) (nvalchemi.hooks.RichLayout method) default_preview_epoch() (nvalchemi.hooks.BaseRichLayout method) (nvalchemi.hooks.DynamicsRichLayout method) (nvalchemi.hooks.RichLayout method) default_preview_history() (nvalchemi.hooks.BaseRichLayout method) (nvalchemi.hooks.DynamicsRichLayout method) (nvalchemi.hooks.RichLayout method) (nvalchemi.hooks.TrainingRichLayout method) default_preview_stage() (nvalchemi.hooks.BaseRichLayout method) (nvalchemi.hooks.DynamicsRichLayout method) (nvalchemi.hooks.RichLayout method) default_training_fn() (in module nvalchemi.training) defrag() (nvalchemi.data.Batch method) defragment() (nvalchemi.data.AtomicDataZarrWriter method) delete() (nvalchemi.data.AtomicDataZarrWriter method) DemoDynamics (class in nvalchemi.dynamics) DemoModel (class in nvalchemi.models.demo) DemoModelWrapper (class in nvalchemi.models.demo) DenseBatchPadder (class in nvalchemi.distributed) DensePadder (class in nvalchemi.distributed) describe() (nvalchemi.distributed.JitAdapter method) (nvalchemi.distributed.MethodAdapter method) (nvalchemi.distributed.OpAdapter method) (nvalchemi.distributed.PythonAdapter method) device (nvalchemi.data.Batch attribute) (nvalchemi.distributed.DistributedManager property) (nvalchemi.dynamics.GPUBuffer attribute) (nvalchemi.dynamics.GPUBuffer property) devices (nvalchemi.training.TrainingStrategy attribute) dftd3() (in module nvalchemiops.torch.interactions.dispersion) DFTD3ModelWrapper (class in nvalchemi.models.dftd3) dipole (nvalchemi.data.AtomicData attribute) direct_derivative_keys() (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) distributed (nvalchemi.distributed.DistributedManager property) distributed_manager (nvalchemi.training.TrainingStrategy attribute) distributed_method() (in module nvalchemi.distributed) distributed_setup() (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) distributed_teardown() (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) DistributedManager (class in nvalchemi.distributed) DistributedModel (class in nvalchemi.distributed) DistributedPipeline (class in nvalchemi.dynamics) DistributedPipelineModel (class in nvalchemi.distributed) distribution_spec() (nvalchemi.models.aimnet2.AIMNet2Wrapper method) (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.dftd3.DFTD3ModelWrapper method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.lj.LennardJonesModelWrapper method) (nvalchemi.models.mace.MACEWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) (nvalchemi.models.uma.UMAWrapper method) DistributionSpec (class in nvalchemi.distributed) domain_widths() (nvalchemi.distributed.SpatialPartitioner method) DomainConfig (class in nvalchemi.distributed) DomainParallel (class in nvalchemi.distributed) drain() (nvalchemi.dynamics.DataSink method), [1] drop_last (nvalchemi.data.DataLoader attribute) dt (nvalchemi.dynamics.DemoDynamics attribute) dtype (nvalchemi.models.demo.DemoModelWrapper property) dtype_policy (nvalchemi.training.BaseLossFunction attribute) (nvalchemi.training.BaseLossFunction property) (nvalchemi.training.ComposedLossFunction attribute) (nvalchemi.training.ComposedLossFunction property) DynamicsContext (class in nvalchemi.hooks) DynamicsRichLayout (class in nvalchemi.hooks) DynamicsStage (class in nvalchemi.dynamics) E edge_embeddings (nvalchemi.data.AtomicData attribute) edge_ptr (nvalchemi.data.Batch property) effective_ghost_width() (nvalchemi.distributed.DomainConfig method) effective_migration_hysteresis() (nvalchemi.distributed.DomainConfig method) effective_read_window (nvalchemi.data.DataLoader property) elapsed_s (nvalchemi.hooks.ScalarSnapshot attribute) EMAHook (class in nvalchemi.training.hooks) embedding_shapes (nvalchemi.models.aimnet2.AIMNet2Wrapper property) (nvalchemi.models.base.BaseModelMixin property) (nvalchemi.models.demo.DemoModelWrapper property) (nvalchemi.models.dftd3.DFTD3ModelWrapper property) (nvalchemi.models.ewald.EwaldModelWrapper property) (nvalchemi.models.lj.LennardJonesModelWrapper property) (nvalchemi.models.mace.MACEWrapper property) (nvalchemi.models.pipeline.PipelineModelWrapper property) (nvalchemi.models.pme.PMEModelWrapper property) (nvalchemi.models.uma.UMAWrapper property) empty() (nvalchemi.data.Batch class method) empty_like() (nvalchemi.data.Batch class method) energy (nvalchemi.data.AtomicData attribute) energy_output (nvalchemi.distributed.CompilePolicy property) EnergyDriftMonitorHook (class in nvalchemi.dynamics.hooks), [1] EnergyHuberLoss (class in nvalchemi.training) EnergyMAELoss (class in nvalchemi.training) EnergyMSELoss (class in nvalchemi.training) entry_status (nvalchemi.dynamics.FusedStage attribute) epoch (nvalchemi.hooks.ScalarSnapshot attribute) (nvalchemi.hooks.TrainContext attribute) epoch_count (nvalchemi.training.TrainingStrategy attribute) epoch_interval (nvalchemi.training.hooks.CheckpointHook attribute) epoch_step_count (nvalchemi.hooks.ScalarSnapshot attribute) (nvalchemi.hooks.TrainContext attribute) (nvalchemi.training.TrainingStrategy attribute) epoch_step_modifier (nvalchemi.training.TrainingStrategy attribute) error (nvalchemi.distributed.AdapterStatus attribute) evaluate() (nvalchemi.dynamics.ConvergenceHook method) event_count (nvalchemi.hooks.ReportingState attribute) (nvalchemi.hooks.ScalarSnapshot attribute) every_n_epochs (nvalchemi.training.ValidationConfig attribute) every_n_steps (nvalchemi.training.ValidationConfig attribute) ewald_real_space() (in module nvalchemiops.torch.interactions.electrostatics.ewald) ewald_reciprocal_space() (in module nvalchemiops.torch.interactions.electrostatics.ewald) EwaldModelWrapper (class in nvalchemi.models.ewald) execute() (nvalchemi.training.ValidationLoop method) exhausted (nvalchemi.dynamics.SizeAwareSampler property) exit_status (nvalchemi.dynamics.BaseDynamics attribute) (nvalchemi.dynamics.FusedStage attribute) export_model() (nvalchemi.models.aimnet2.AIMNet2Wrapper method) (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.demo.DemoModelWrapper method) (nvalchemi.models.dftd3.DFTD3ModelWrapper method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.lj.LennardJonesModelWrapper method) (nvalchemi.models.mace.MACEWrapper method) (nvalchemi.models.pipeline.PipelineModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) extra_keys (nvalchemi.dynamics.hooks.NaNDetectorHook attribute), [1] extra_repr() (nvalchemi.models.aimnet2.AIMNet2Wrapper method) (nvalchemi.models.demo.DemoModelWrapper method) (nvalchemi.models.dftd3.DFTD3ModelWrapper method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.lj.LennardJonesModelWrapper method) (nvalchemi.models.mace.MACEWrapper method) (nvalchemi.models.pipeline.PipelineModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) (nvalchemi.models.uma.UMAWrapper method) (nvalchemi.training.ComposedLossFunction method) (nvalchemi.training.EnergyHuberLoss method) (nvalchemi.training.EnergyMAELoss method) (nvalchemi.training.EnergyMSELoss method) (nvalchemi.training.ForceHuberLoss method) (nvalchemi.training.ForceL2NormLoss method) (nvalchemi.training.ForceMSELoss method) (nvalchemi.training.StressHuberLoss method) (nvalchemi.training.StressMSELoss method) F field_levels (nvalchemi.data.AtomicDataZarrReader property) (nvalchemi.data.Reader property) field_names (nvalchemi.data.datapipes.MultiDataset property) (nvalchemi.data.Dataset property) (nvalchemi.data.InMemoryDataset property) (nvalchemi.data.Reader property) field_overrides (nvalchemi.data.datapipes.ZarrWriteConfig attribute) filters (nvalchemi.data.datapipes.ZarrArrayConfig attribute) find_unused_parameters (nvalchemi.distributed.DistributedManager property) (nvalchemi.training.hooks.DDPHook attribute) fire2_step_coord() (in module nvalchemiops.torch.fire2) fire2_step_coord_cell() (in module nvalchemiops.torch.fire2) fire_step() (in module nvalchemiops.dynamics.optimizers.fire) fire_update() (in module nvalchemiops.dynamics.optimizers.fire) ForceHuberLoss (class in nvalchemi.training) ForceL2NormLoss (class in nvalchemi.training) ForceMSELoss (class in nvalchemi.training) forces (nvalchemi.data.AtomicData attribute) forces_via_autograd (nvalchemi.distributed.CompilePolicy property) ForceStrategy (class in nvalchemi.distributed) format (nvalchemi.models.base.NeighborConfig attribute) forward() (nvalchemi.models.aimnet2.AIMNet2Wrapper method) (nvalchemi.models.demo.DemoModel method) (nvalchemi.models.demo.DemoModelWrapper method) (nvalchemi.models.dftd3.DFTD3ModelWrapper method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.lj.LennardJonesModelWrapper method) (nvalchemi.models.mace.MACEWrapper method) (nvalchemi.models.pipeline.PipelineModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) (nvalchemi.models.uma.UMAWrapper method) (nvalchemi.training.BaseLossFunction method) (nvalchemi.training.ComposedLossFunction method) freeze_category (nvalchemi.dynamics.hooks.FreezeAtomsHook attribute), [1] freeze_mode (nvalchemi.training.FineTuningStrategy attribute) (nvalchemi.training.hooks.TrainableParameterHook attribute) freeze_patterns (nvalchemi.training.FineTuningStrategy attribute) (nvalchemi.training.hooks.TrainableParameterHook attribute) FreezeAtomsHook (class in nvalchemi.dynamics.hooks), [1] frequency (nvalchemi.dynamics.ConvergenceHook attribute) (nvalchemi.dynamics.hooks.EnergyDriftMonitorHook attribute), [1] (nvalchemi.dynamics.hooks.FreezeAtomsHook attribute), [1] (nvalchemi.dynamics.hooks.MaxForceClampHook attribute), [1] (nvalchemi.dynamics.hooks.NaNDetectorHook attribute), [1] (nvalchemi.dynamics.hooks.SnapshotHook attribute), [1] (nvalchemi.dynamics.hooks.StageTimingHook attribute), [1] (nvalchemi.hooks.BiasedPotentialHook attribute) (nvalchemi.hooks.Hook attribute) (nvalchemi.hooks.StageTimingHook attribute) (nvalchemi.hooks.TorchProfilerHook attribute) (nvalchemi.hooks.WrapPeriodicHook attribute) (nvalchemi.training.hooks.TrainingUpdateOrchestrator attribute) frobenius_mse() (in module nvalchemi.training.losses.reductions) from_batch() (nvalchemi.distributed.DistributedModel method) (nvalchemi.distributed.ShardedBatch static method) from_checkpoint() (nvalchemi.models.aimnet2.AIMNet2Wrapper class method) (nvalchemi.models.demo.DemoModelWrapper class method) (nvalchemi.models.mace.MACEWrapper class method) (nvalchemi.models.uma.UMAWrapper class method) from_data_list() (nvalchemi.data.Batch class method) from_dict() (nvalchemi.distributed.DistributionSpec class method) (nvalchemi.distributed.JitAdapter class method) (nvalchemi.distributed.MethodAdapter class method) (nvalchemi.distributed.MLIPSpec class method) (nvalchemi.distributed.OpAdapter class method) (nvalchemi.distributed.PythonAdapter class method) from_fmax() (nvalchemi.dynamics.ConvergenceHook class method) from_forces() (nvalchemi.dynamics.ConvergenceHook class method) from_pretrained_checkpoint() (nvalchemi.training.FineTuningStrategy class method) from_raw_dicts() (nvalchemi.data.Batch class method) from_training_strategy() (nvalchemi.training.ValidationLoop class method) full_batch() (nvalchemi.distributed.ShardedBatch method) FunctionAdapter (class in nvalchemi.distributed) FusedStage (class in nvalchemi.dynamics) G gather() (nvalchemi.distributed.DomainParallel method) get_available_backend() (nvalchemi.distributed.DistributedManager static method) get_averaged_model() (nvalchemi.training.hooks.EMAHook method) get_batch() (nvalchemi.data.Dataset method) get_data() (nvalchemi.data.Batch method) get_fused_batches() (nvalchemi.data.datapipes.MultiDataset method) (nvalchemi.data.Dataset method) (nvalchemi.data.InMemoryDataset method) get_mesh_group() (nvalchemi.distributed.DistributedManager method) get_metadata() (nvalchemi.data.AtomicDataZarrReader method) (nvalchemi.data.datapipes.MultiDataset method) (nvalchemi.data.Dataset method) (nvalchemi.data.InMemoryDataset method) get_neighbor_ranks() (nvalchemi.distributed.SpatialPartitioner method) GLOBAL (nvalchemi.distributed.HookScope attribute) global_mesh (nvalchemi.distributed.DistributedManager property) global_rank (nvalchemi.dynamics.DataSink property) (nvalchemi.dynamics.DistributedPipeline property) (nvalchemi.hooks.HookContext attribute) (nvalchemi.hooks.ReportingOrchestrator property) (nvalchemi.hooks.ScalarSnapshot attribute) global_step_count (nvalchemi.hooks.TrainContext attribute) (nvalchemi.training.TrainingStrategy attribute) GPUBuffer (class in nvalchemi.dynamics) grad_mode (nvalchemi.training.ValidationConfig attribute) grad_scaler (nvalchemi.hooks.TrainContext attribute) gradient_keys (nvalchemi.models.base.ModelConfig attribute) gradients (nvalchemi.hooks.TrainContext attribute) graph_alpha_spins (nvalchemi.data.AtomicData attribute) graph_embeddings (nvalchemi.data.AtomicData attribute) GRAPH_PARTITION (nvalchemi.distributed.config.StrategyKind attribute) GraphPadder (class in nvalchemi.distributed) group() (nvalchemi.distributed.DistributedManager method) group_name() (nvalchemi.distributed.DistributedManager method) group_names (nvalchemi.distributed.DistributedManager property) group_rank() (nvalchemi.distributed.DistributedManager method) group_size() (nvalchemi.distributed.DistributedManager method) H half_list (nvalchemi.models.base.NeighborConfig attribute) HALO (nvalchemi.distributed.config.StrategyKind attribute) has_pending_fused_batches() (nvalchemi.data.datapipes.MultiDataset method) (nvalchemi.data.Dataset method) (nvalchemi.data.InMemoryDataset method) history (nvalchemi.hooks.RichReporter property) Hook (class in nvalchemi.hooks) HookContext (class in nvalchemi.hooks) HookRegistryMixin (class in nvalchemi.hooks) hooks (nvalchemi.dynamics.BaseDynamics attribute) (nvalchemi.dynamics.DemoDynamics attribute) (nvalchemi.hooks.HookRegistryMixin attribute) (nvalchemi.training.TrainingStrategy attribute) HookScope (class in nvalchemi.distributed) HostMemory (class in nvalchemi.dynamics) I in_memory_batch (nvalchemi.data.InMemoryDataset attribute) include_dynamics_scalars (nvalchemi.hooks.BaseRichLayout attribute) include_kinetic (nvalchemi.dynamics.hooks.EnergyDriftMonitorHook attribute), [1] index_select() (nvalchemi.data.Batch method) inference_model (nvalchemi.training.TrainingStrategy attribute) info (nvalchemi.data.AtomicData attribute) init_distributed() (nvalchemi.dynamics.DistributedPipeline method) initialize() (nvalchemi.distributed.DistributedManager static method) initialize_env() (nvalchemi.distributed.DistributedManager static method) initialize_mesh() (nvalchemi.distributed.DistributedManager method) initialize_open_mpi() (nvalchemi.distributed.DistributedManager static method) initialize_slurm() (nvalchemi.distributed.DistributedManager static method) InMemoryDataset (class in nvalchemi.data) input_data() (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.dftd3.DFTD3ModelWrapper method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) install() (nvalchemi.distributed.AdapterRegistry method) (nvalchemi.distributed.JitAdapter method) (nvalchemi.distributed.MethodAdapter method) (nvalchemi.distributed.OpAdapter method) (nvalchemi.distributed.PythonAdapter method) install_site (nvalchemi.distributed.AdapterStatus attribute) invalidate_cache() (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) irecv() (nvalchemi.data.Batch class method) is_full (nvalchemi.dynamics.DataSink property) is_initialized() (nvalchemi.distributed.DistributedManager class method) is_rank_zero (nvalchemi.hooks.ReportingOrchestrator property) isend() (nvalchemi.data.Batch method) iter_hooks() (nvalchemi.training.hooks.TrainingUpdateOrchestrator method) J JitAdapter (class in nvalchemi.distributed) K keeps_owner() (nvalchemi.distributed.SpatialPartitioner method) keys (nvalchemi.data.Batch attribute) kind (nvalchemi.distributed.AdapterStatus attribute) kinetic_energies (nvalchemi.data.AtomicData attribute) L langevin_baoab_finalize() (in module nvalchemiops.dynamics.integrators) langevin_baoab_half_step() (in module nvalchemiops.dynamics.integrators) last_checkpoint_index (nvalchemi.training.hooks.CheckpointHook attribute) last_event_at_s (nvalchemi.hooks.ReportingState attribute) last_global_rank (nvalchemi.hooks.ReportingState attribute) last_stage (nvalchemi.hooks.ReportingState attribute) last_step_count (nvalchemi.hooks.ReportingState attribute) last_validation (nvalchemi.training.TrainingStrategy attribute) LennardJonesModelWrapper (class in nvalchemi.models.lj) list_active() (nvalchemi.distributed.AdapterRegistry method) lj_energy_forces() (in module nvalchemiops.interactions.lj) lj_energy_forces_virial() (in module nvalchemiops.interactions.lj) load() (nvalchemi.distributed.MLIPSpec class method) (nvalchemi.models.pipeline.PipelineModelWrapper class method) load_batches() (nvalchemi.data.datapipes.MultiDataset method) (nvalchemi.data.Dataset method) (nvalchemi.data.InMemoryDataset method) load_checkpoint() (in module nvalchemi.training) (nvalchemi.training.FineTuningStrategy class method) (nvalchemi.training.TrainingStrategy class method) load_state_dict() (nvalchemi.hooks.CheckpointableHook method) (nvalchemi.training.hooks.EMAHook method) LOCAL (nvalchemi.distributed.HookScope attribute) (nvalchemi.distributed.Scope attribute) local_batch (nvalchemi.distributed.ShardedBatch property) local_batch_with_edges() (nvalchemi.distributed.ShardedBatch method) local_rank (nvalchemi.distributed.DistributedManager property) (nvalchemi.dynamics.DataSink property) (nvalchemi.dynamics.DistributedPipeline property) local_stage (nvalchemi.dynamics.DistributedPipeline property) localize() (in module nvalchemi.distributed) LoggingHook (class in nvalchemi.dynamics.hooks), [1] loss (nvalchemi.hooks.TrainContext attribute) loss_fn (nvalchemi.training.TrainingStrategy attribute) (nvalchemi.training.ValidationConfig attribute) loss_target_assembler (nvalchemi.training.TrainingStrategy attribute) losses (nvalchemi.hooks.TrainContext attribute) LossWeightSchedule (class in nvalchemi.training) lr_schedulers (nvalchemi.hooks.TrainContext attribute) M MACEWrapper (class in nvalchemi.models.mace) make_neighbor_hooks() (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.pipeline.PipelineModelWrapper method) mark_event() (nvalchemi.hooks.ReportingState method) mask() (nvalchemi.training.BaseLossFunction method) (nvalchemi.training.EnergyHuberLoss method) (nvalchemi.training.EnergyMAELoss method) (nvalchemi.training.EnergyMSELoss method) (nvalchemi.training.ForceL2NormLoss method) (nvalchemi.training.ForceMSELoss method) (nvalchemi.training.StressMSELoss method) masked_update() (nvalchemi.dynamics.BaseDynamics method) max_atoms (nvalchemi.dynamics.SizeAwareSampler property) max_batch_size (nvalchemi.dynamics.SizeAwareSampler property) max_edges (nvalchemi.dynamics.SizeAwareSampler property) max_force (nvalchemi.dynamics.hooks.MaxForceClampHook attribute), [1] max_messages (nvalchemi.hooks.ReportingState attribute) max_num_nodes (nvalchemi.data.Batch property) MaxForceClampHook (class in nvalchemi.dynamics.hooks), [1] merge() (nvalchemi.distributed.MLIPSpec method) mesh() (nvalchemi.distributed.DistributedManager method) mesh_dims (nvalchemi.distributed.DistributedManager property) mesh_names() (nvalchemi.distributed.DistributedManager method) mesh_sizes() (nvalchemi.distributed.DistributedManager method) messages (nvalchemi.hooks.ReportingState attribute) (nvalchemi.hooks.ScalarSnapshot attribute) meta (nvalchemi.data.datapipes.ZarrWriteConfig attribute) metadata (nvalchemi.hooks.ReportingState attribute) MethodAdapter (class in nvalchemi.distributed) metric (nvalchemi.dynamics.hooks.EnergyDriftMonitorHook attribute), [1] MLIPSpec (class in nvalchemi.distributed) model (nvalchemi.dynamics.BaseDynamics attribute) (nvalchemi.dynamics.DemoDynamics attribute) (nvalchemi.hooks.HookContext attribute) (nvalchemi.models.aimnet2.AIMNet2Wrapper attribute) (nvalchemi.models.mace.MACEWrapper attribute) model_config (nvalchemi.distributed.DomainConfig attribute) (nvalchemi.models.aimnet2.AIMNet2Wrapper attribute) (nvalchemi.models.dftd3.DFTD3ModelWrapper attribute) (nvalchemi.models.ewald.EwaldModelWrapper attribute) (nvalchemi.models.lj.LennardJonesModelWrapper attribute) (nvalchemi.models.mace.MACEWrapper attribute) (nvalchemi.models.pipeline.PipelineModelWrapper attribute) (nvalchemi.models.pme.PMEModelWrapper attribute) (nvalchemi.models.uma.UMAWrapper attribute) (nvalchemi.training.hooks.EMAHook attribute) model_dump() (nvalchemi.data.Batch method) model_is_conservative (nvalchemi.dynamics.BaseDynamics attribute) (nvalchemi.dynamics.BaseDynamics property) model_keys (nvalchemi.training.hooks.DDPHook attribute) model_post_init() (nvalchemi.training.hooks.EMAHook method) models (nvalchemi.hooks.TrainContext attribute) (nvalchemi.training.CheckpointManifest attribute) (nvalchemi.training.TrainingStrategy attribute) modify_ema_methods() (nvalchemi.models.mace.MACEWrapper method) module_patches (nvalchemi.training.FineTuningStrategy attribute) momenta (nvalchemi.data.AtomicData attribute) MultiDataset (class in nvalchemi.data.datapipes) MultiDatasetBatchSampler (class in nvalchemi.data.datapipes) MultiDatasetSampler (class in nvalchemi.data.datapipes) N n_global (nvalchemi.distributed.ShardedBatch property) n_owned (nvalchemi.distributed.ShardedBatch property) n_steps (nvalchemi.dynamics.BaseDynamics attribute) name (nvalchemi.hooks.BaseRichLayout attribute) (nvalchemi.hooks.TorchProfilerHook attribute) (nvalchemi.training.ValidationConfig attribute) NaNDetectorHook (class in nvalchemi.dynamics.hooks), [1] needs_pbc (nvalchemi.models.base.ModelConfig attribute) neighbor_config (nvalchemi.models.base.ModelConfig attribute) neighbor_list (nvalchemi.data.AtomicData attribute) neighbor_list_shifts (nvalchemi.data.AtomicData attribute) neighbor_matrix (nvalchemi.data.AtomicData attribute) neighbor_matrix_shifts (nvalchemi.data.AtomicData attribute) neighbor_ranks (nvalchemi.distributed.ParticleHaloConfig property) neighbor_refresh_adapters() (in module nvalchemi.distributed) neighbor_span() (nvalchemi.distributed.SpatialPartitioner method) NeighborListHook (class in nvalchemi.hooks) next_rank (nvalchemi.dynamics.DemoDynamics attribute) nhc_compute_2ke() (in module nvalchemiops.dynamics.integrators) nhc_compute_chain_energy() (in module nvalchemiops.dynamics.integrators) nhc_compute_masses() (in module nvalchemiops.dynamics.integrators) nhc_position_update() (in module nvalchemiops.dynamics.integrators) nhc_thermostat_chain_update() (in module nvalchemiops.dynamics.integrators) nhc_velocity_half_step() (in module nvalchemiops.dynamics.integrators) node_alpha_spins (nvalchemi.data.AtomicData attribute) node_attrs (nvalchemi.data.AtomicData attribute) node_beta_spins (nvalchemi.data.AtomicData attribute) node_embeddings (nvalchemi.data.AtomicData attribute) normalize() (nvalchemi.training.BaseLossFunction method) (nvalchemi.training.EnergyHuberLoss method) (nvalchemi.training.EnergyMAELoss method) (nvalchemi.training.EnergyMSELoss method) normalize_weights (nvalchemi.training.ComposedLossFunction attribute) nph_barostat_half_step() (in module nvalchemiops.dynamics.integrators) nph_velocity_half_step() (in module nvalchemiops.dynamics.integrators) npt_barostat_half_step() (in module nvalchemiops.dynamics.integrators) npt_cell_update() (in module nvalchemiops.dynamics.integrators) npt_position_update() (in module nvalchemiops.dynamics.integrators) npt_thermostat_half_step() (in module nvalchemiops.dynamics.integrators) npt_velocity_half_step() (in module nvalchemiops.dynamics.integrators) num_criteria (nvalchemi.dynamics.ConvergenceHook property) num_edges (nvalchemi.data.Batch property) (nvalchemi.dynamics.base.BufferConfig attribute) num_edges_list (nvalchemi.data.Batch property) num_edges_per_graph (nvalchemi.data.Batch property) num_epochs (nvalchemi.training.TrainingStrategy attribute) num_graphs (nvalchemi.data.Batch property) (nvalchemi.distributed.ShardedBatch property) num_neighbors (nvalchemi.data.AtomicData attribute) num_nodes (nvalchemi.data.Batch property) (nvalchemi.dynamics.base.BufferConfig attribute) num_nodes_list (nvalchemi.data.Batch property) num_nodes_per_graph (nvalchemi.data.Batch property) num_steps (nvalchemi.training.TrainingStrategy attribute) num_streams (nvalchemi.data.DataLoader attribute) num_systems (nvalchemi.dynamics.base.BufferConfig attribute) num_workers (nvalchemi.data.Dataset attribute) O offsets (nvalchemi.data.datapipes.MultiDataset property) ON_CONVERGE (nvalchemi.dynamics.DynamicsStage attribute) on_register (nvalchemi.hooks.Hook attribute) on_trace_ready_path (nvalchemi.hooks.TorchProfilerHook attribute) OpAdapter (class in nvalchemi.distributed) optimizer_cls (nvalchemi.training.OptimizerConfig attribute) optimizer_configs (nvalchemi.training.TrainingStrategy attribute) optimizer_kwargs (nvalchemi.training.OptimizerConfig attribute) optimizer_step_skipped (nvalchemi.training.hooks.TrainingUpdateOrchestrator property) optimizers (nvalchemi.hooks.TrainContext attribute) (nvalchemi.training.CheckpointManifest attribute) optional_inputs (nvalchemi.models.base.ModelConfig attribute) output_data() (nvalchemi.models.base.BaseModelMixin method) (nvalchemi.models.dftd3.DFTD3ModelWrapper method) (nvalchemi.models.ewald.EwaldModelWrapper method) (nvalchemi.models.lj.LennardJonesModelWrapper method) (nvalchemi.models.pme.PMEModelWrapper method) output_dir (nvalchemi.hooks.TorchProfilerHook attribute) OutputKind (class in nvalchemi.distributed) outputs (nvalchemi.models.base.ModelConfig attribute) OutputSpec (class in nvalchemi.distributed) OWNED (nvalchemi.distributed.Scope attribute) P pad() (nvalchemi.distributed.COOPadder method) (nvalchemi.distributed.DenseBatchPadder method) (nvalchemi.distributed.DensePadder method) (nvalchemi.distributed.GraphPadder method) ParticleHaloConfig (class in nvalchemi.distributed) partition() (nvalchemi.distributed.DomainParallel method) partition_mode (nvalchemi.distributed.ShardedBatch property) patches (nvalchemi.training.hooks.ModulePatchHook attribute) pbc (nvalchemi.data.AtomicData attribute) pbc_shifts (nvalchemi.distributed.ParticleHaloConfig property) per_epoch (nvalchemi.training.LossWeightSchedule attribute) per_graph_mean() (in module nvalchemi.training.losses.reductions) per_graph_sum() (in module nvalchemi.training.losses.reductions) per_sample_loss (nvalchemi.training.BaseLossFunction attribute) PhysicsNeMoUninitializedDistributedManagerWarning (class in nvalchemi.distributed) pin_memory (nvalchemi.data.DataLoader attribute) (nvalchemi.data.Dataset property) (nvalchemi.data.InMemoryDataset property) pin_memory() (nvalchemi.data.Batch method) PipelineGroup (class in nvalchemi.models.pipeline) PipelineModelWrapper (class in nvalchemi.models.pipeline) PipelineStep (class in nvalchemi.models.pipeline) PMEModelWrapper (class in nvalchemi.models.pme) positions (nvalchemi.data.AtomicData attribute) post_update() (nvalchemi.dynamics.BaseDynamics method) (nvalchemi.dynamics.DemoDynamics method) pre_update() (nvalchemi.dynamics.BaseDynamics method) (nvalchemi.dynamics.DemoDynamics method) precision (nvalchemi.training.hooks.MixedPrecisionHook attribute) prefetch() (nvalchemi.data.datapipes.MultiDataset method) (nvalchemi.data.Dataset method) (nvalchemi.data.InMemoryDataset method) prefetch_batch() (nvalchemi.data.datapipes.MultiDataset method) (nvalchemi.data.Dataset method) prefetch_count (nvalchemi.data.datapipes.MultiDataset property) (nvalchemi.data.Dataset property) (nvalchemi.data.InMemoryDataset property) prefetch_factor (nvalchemi.data.DataLoader attribute) prefetch_fused_batches() (nvalchemi.data.datapipes.MultiDataset method) (nvalchemi.data.Dataset method) (nvalchemi.data.InMemoryDataset method) prefetch_many() (nvalchemi.data.Dataset method) preview() (nvalchemi.hooks.RichReporter class method) prior_rank (nvalchemi.dynamics.DemoDynamics attribute) priority (nvalchemi.training.hooks.TrainingUpdateHook attribute) process_group (nvalchemi.training.hooks.DDPHook attribute) profile_memory (nvalchemi.hooks.TorchProfilerHook attribute) put() (nvalchemi.data.Batch method) PythonAdapter (class in nvalchemi.distributed) R rank (nvalchemi.distributed.DistributedManager property) rank_assignment (nvalchemi.distributed.ShardedBatch property) rank_subdirs (nvalchemi.hooks.TorchProfilerHook attribute) rank_to_cell_bounds() (nvalchemi.distributed.SpatialPartitioner method) rank_to_grid_coords() (nvalchemi.distributed.SpatialPartitioner method) RANK_ZERO (nvalchemi.distributed.HookScope attribute) rank_zero_only (nvalchemi.training.hooks.CheckpointHook attribute) read() (nvalchemi.data.Reader method) (nvalchemi.dynamics.DataSink method), [1] (nvalchemi.dynamics.GPUBuffer method) (nvalchemi.dynamics.HostMemory method) (nvalchemi.dynamics.ZarrData method) read_many() (nvalchemi.data.Dataset method) (nvalchemi.data.Reader method) Reader (class in nvalchemi.data) reader (nvalchemi.data.Dataset attribute) record_shapes (nvalchemi.hooks.TorchProfilerHook attribute) recv() (nvalchemi.data.Batch class method) Reduce (class in nvalchemi.distributed) reduce() (nvalchemi.training.BaseLossFunction method) (nvalchemi.training.ForceL2NormLoss method) (nvalchemi.training.ForceMSELoss method) (nvalchemi.training.StressMSELoss method) ReductionContext (class in nvalchemi.training) refill_check() (nvalchemi.dynamics.BaseDynamics method) (nvalchemi.dynamics.FusedStage method) refine_grid_for_ranks() (nvalchemi.distributed.SpatialPartitioner static method) refresh() (nvalchemi.data.AtomicDataZarrReader method) refresh_neighbors() (in module nvalchemi.distributed) register_bookkeeping_key() (nvalchemi.dynamics.BaseDynamics class method) register_fused_hook() (nvalchemi.dynamics.FusedStage method) register_hook() (nvalchemi.hooks.HookRegistryMixin method) register_type_serializer() (in module nvalchemi.training) render() (nvalchemi.hooks.BaseRichLayout method) (nvalchemi.hooks.DynamicsRichLayout method) (nvalchemi.hooks.RichLayout method) (nvalchemi.hooks.TrainingRichLayout method) renderable() (nvalchemi.hooks.RichReporter method) report() (nvalchemi.hooks.Reporter method) (nvalchemi.hooks.RichReporter method) (nvalchemi.hooks.TensorBoardReporter method) Reporter (class in nvalchemi.hooks) ReportingOrchestrator (class in nvalchemi.hooks) ReportingState (class in nvalchemi.hooks) request_replacement() (nvalchemi.dynamics.SizeAwareSampler method) request_replacements() (nvalchemi.dynamics.SizeAwareSampler method) request_replacements_budget() (nvalchemi.dynamics.SizeAwareSampler method) required_inputs (nvalchemi.models.base.ModelConfig attribute) requires_eval_grad (nvalchemi.training.BaseLossFunction attribute) requires_eval_grad() (nvalchemi.training.ComposedLossFunction method) reset() (nvalchemi.hooks.StageTimingHook method) reshard_by_destination() (in module nvalchemi.distributed) resolve_cap() (in module nvalchemi.distributed) resolve_global_rank() (in module nvalchemi.distributed) resolve_world_size() (in module nvalchemi.distributed) restore() (nvalchemi.distributed.AdapterRegistry method) (nvalchemi.distributed.JitAdapter method) (nvalchemi.distributed.MethodAdapter method) (nvalchemi.distributed.OpAdapter method) (nvalchemi.distributed.PythonAdapter method) restore_checkpoint() (nvalchemi.training.TrainingStrategy method) RichLayout (class in nvalchemi.hooks) RichReporter (class in nvalchemi.hooks) run() (nvalchemi.distributed.DomainParallel method) (nvalchemi.dynamics.BaseDynamics method) (nvalchemi.dynamics.DistributedPipeline method) (nvalchemi.dynamics.FusedStage method) S sampler (nvalchemi.data.DataLoader attribute) sampler_cls (nvalchemi.training.hooks.DDPHook attribute) sampler_kwargs (nvalchemi.training.hooks.DDPHook attribute) save() (nvalchemi.distributed.MLIPSpec method) (nvalchemi.models.pipeline.PipelineModelWrapper method) save_checkpoint() (in module nvalchemi.training) (nvalchemi.training.TrainingStrategy method) scalars (nvalchemi.hooks.ScalarSnapshot attribute) ScalarSnapshot (class in nvalchemi.hooks) scatter_to_owners() (in module nvalchemi.distributed) schedule (nvalchemi.hooks.TorchProfilerHook attribute) scheduler_cls (nvalchemi.training.OptimizerConfig attribute) scheduler_kwargs (nvalchemi.training.OptimizerConfig attribute) scheduler_metric_adapter (nvalchemi.training.OptimizerConfig attribute) schedulers (nvalchemi.training.CheckpointManifest attribute) schema_version (nvalchemi.training.CheckpointManifest attribute) Scope (class in nvalchemi.distributed) seed_history() (nvalchemi.hooks.RichReporter method) send() (nvalchemi.data.Batch method) serializer (nvalchemi.data.datapipes.ZarrArrayConfig attribute) set_config() (nvalchemi.models.base.BaseModelMixin method) set_epoch() (nvalchemi.data.DataLoader method) (nvalchemi.data.datapipes.MultiDatasetBatchSampler method) (nvalchemi.data.datapipes.MultiDatasetSampler method) set_epoch_step() (nvalchemi.data.DataLoader method) set_eval (nvalchemi.training.ValidationConfig attribute) setup() (nvalchemi.distributed.DistributedManager static method) (nvalchemi.dynamics.DistributedPipeline method) setup_optimizers() (in module nvalchemi.training) shard_size (nvalchemi.data.datapipes.ZarrArrayConfig attribute) ShardedBatch (class in nvalchemi.distributed) shifts (nvalchemi.data.AtomicData attribute) single_model_input (nvalchemi.training.TrainingStrategy attribute) sink (nvalchemi.dynamics.hooks.SnapshotHook attribute), [1] SizeAwareSampler (class in nvalchemi.dynamics) skin (nvalchemi.models.base.NeighborConfig attribute) SnapshotHook (class in nvalchemi.dynamics.hooks), [1] source_status (nvalchemi.dynamics.ConvergenceHook attribute) SpatialPartitioner (class in nvalchemi.distributed) SPEC_DFTD3_HALO (in module nvalchemi.distributed.spec) SPEC_EWALD_HALO (in module nvalchemi.distributed.spec) SPEC_LJ_HALO (in module nvalchemi.distributed.spec) SPEC_MPNN_GP (in module nvalchemi.distributed.spec) SPEC_MPNN_HALO (in module nvalchemi.distributed.spec) SPEC_PME_HALO (in module nvalchemi.distributed.spec) SPEC_UMA_HALO (in module nvalchemi.distributed.spec) spin (nvalchemi.data.AtomicData attribute) spline_gather_with_force() (in module nvalchemiops.torch.spline) spline_spread() (in module nvalchemiops.torch.spline) stage (nvalchemi.dynamics.ConvergenceHook attribute) (nvalchemi.dynamics.hooks.EnergyDriftMonitorHook attribute), [1] (nvalchemi.dynamics.hooks.FreezeAtomsHook attribute), [1] (nvalchemi.dynamics.hooks.MaxForceClampHook attribute), [1] (nvalchemi.dynamics.hooks.NaNDetectorHook attribute), [1] (nvalchemi.dynamics.hooks.SnapshotHook attribute), [1] (nvalchemi.hooks.BiasedPotentialHook attribute) (nvalchemi.hooks.Hook attribute) (nvalchemi.hooks.ScalarSnapshot attribute) (nvalchemi.hooks.WrapPeriodicHook attribute) (nvalchemi.training.hooks.TrainingUpdateOrchestrator attribute) stages (nvalchemi.dynamics.DistributedPipeline attribute) StageTimingHook (class in nvalchemi.dynamics.hooks), [1] (class in nvalchemi.hooks) started_at_s (nvalchemi.hooks.ReportingState attribute) state (nvalchemi.distributed.AdapterStatus attribute) state_dict() (nvalchemi.hooks.CheckpointableHook method) (nvalchemi.training.hooks.EMAHook method) static_graph (nvalchemi.training.hooks.DDPHook attribute) step() (nvalchemi.distributed.DomainParallel method) (nvalchemi.dynamics.BaseDynamics method) (nvalchemi.dynamics.DistributedPipeline method) (nvalchemi.dynamics.FusedStage method) step_count (nvalchemi.dynamics.BaseDynamics attribute) (nvalchemi.dynamics.DemoDynamics attribute) (nvalchemi.hooks.DynamicsContext attribute) (nvalchemi.hooks.HookRegistryMixin attribute) (nvalchemi.hooks.ScalarSnapshot attribute) (nvalchemi.hooks.TrainContext attribute) (nvalchemi.training.TrainingStrategy attribute) step_interval (nvalchemi.training.hooks.CheckpointHook attribute) step_lr_schedulers() (in module nvalchemi.training) step_optimizers() (in module nvalchemi.training) store (nvalchemi.dynamics.ZarrData attribute) StrategyKind (class in nvalchemi.distributed.config) stress (nvalchemi.data.AtomicData attribute) StressHuberLoss (class in nvalchemi.training) StressMSELoss (class in nvalchemi.training) sub_stages (nvalchemi.dynamics.FusedStage attribute) summary() (nvalchemi.hooks.StageTimingHook method) supports_pbc (nvalchemi.models.base.ModelConfig attribute) synchronized (nvalchemi.dynamics.DistributedPipeline attribute) system_capacity (nvalchemi.data.Batch property) system_sum() (in module nvalchemi.distributed) T target (nvalchemi.distributed.AdapterStatus attribute) target_device (nvalchemi.data.Dataset attribute) (nvalchemi.data.InMemoryDataset attribute) target_status (nvalchemi.dynamics.ConvergenceHook attribute) task_name (nvalchemi.models.uma.UMAWrapper attribute) TensorBoardReporter (class in nvalchemi.hooks) threshold (nvalchemi.dynamics.hooks.EnergyDriftMonitorHook attribute), [1] timestamp (nvalchemi.training.BaseSpec attribute) timestamp_s (nvalchemi.hooks.ScalarSnapshot attribute) timings (nvalchemi.dynamics.hooks.StageTimingHook attribute), [1] (nvalchemi.hooks.StageTimingHook attribute) to() (nvalchemi.data.Batch method) to_data_list() (nvalchemi.data.Batch method) to_dict() (nvalchemi.distributed.DistributionSpec method) (nvalchemi.distributed.JitAdapter method) (nvalchemi.distributed.MethodAdapter method) (nvalchemi.distributed.MLIPSpec method) (nvalchemi.distributed.OpAdapter method) (nvalchemi.distributed.PythonAdapter method) to_global_batch() (nvalchemi.distributed.ShardedBatch method) to_global_index() (nvalchemi.data.datapipes.MultiDataset method) to_local() (in module nvalchemi.distributed) to_local_index() (nvalchemi.data.datapipes.MultiDataset method) to_spec() (nvalchemi.training.LossWeightSchedule method) topology_version (nvalchemi.distributed.SpatialPartitioner attribute) TorchProfilerHook (class in nvalchemi.dynamics.hooks), [1] (class in nvalchemi.training.hooks) trace_and_validate() (in module nvalchemi.distributed) trainable_patterns (nvalchemi.training.FineTuningStrategy attribute) (nvalchemi.training.hooks.TrainableParameterHook attribute) TrainContext (class in nvalchemi.hooks) training_fn (nvalchemi.training.TrainingStrategy attribute) TrainingRichLayout (class in nvalchemi.hooks) TrainingUpdateHook (class in nvalchemi.training.hooks) TrainingUpdateOrchestrator (class in nvalchemi.training.hooks) trim() (nvalchemi.data.Batch method) U UMAWrapper (class in nvalchemi.models.uma) unpad() (nvalchemi.distributed.COOPadder method) (nvalchemi.distributed.DenseBatchPadder method) (nvalchemi.distributed.DensePadder method) (nvalchemi.distributed.GraphPadder method) update_cell() (nvalchemi.distributed.SpatialPartitioner method) update_from_batch() (nvalchemi.distributed.ShardedBatch method) use_ema (nvalchemi.training.ValidationConfig attribute) use_mixed_precision (nvalchemi.training.ValidationConfig attribute) use_streams (nvalchemi.data.DataLoader attribute) V validate() (nvalchemi.training.BaseLossFunction method) validate_field_names() (nvalchemi.data.datapipes.MultiDataset method) validation (nvalchemi.hooks.TrainContext attribute) validation_config (nvalchemi.training.TrainingStrategy attribute) validation_data (nvalchemi.training.ValidationConfig attribute) validation_fn (nvalchemi.training.ValidationConfig attribute) ValidationLoop (class in nvalchemi.training) value (nvalchemi.training.ConstantWeight attribute) values (nvalchemi.training.PiecewiseWeight attribute) velocities (nvalchemi.data.AtomicData attribute) velocity_rescale() (in module nvalchemiops.dynamics.integrators) velocity_verlet_position_update() (in module nvalchemiops.dynamics.integrators) velocity_verlet_velocity_finalize() (in module nvalchemiops.dynamics.integrators) virial (nvalchemi.data.AtomicData attribute) W weight_factors() (nvalchemi.training.ComposedLossFunction method) with_adapters() (nvalchemi.distributed.MLIPSpec method) with_compile() (nvalchemi.distributed.MLIPSpec method) with_flops (nvalchemi.hooks.TorchProfilerHook attribute) with_outputs() (nvalchemi.distributed.MLIPSpec method) with_stack (nvalchemi.hooks.TorchProfilerHook attribute) workflow (nvalchemi.hooks.HookContext attribute) world_size (nvalchemi.distributed.DistributedManager property) wrapper (nvalchemi.distributed.DistributedModel property) WrapPeriodicHook (class in nvalchemi.hooks) write() (nvalchemi.data.AtomicDataZarrWriter method) (nvalchemi.dynamics.DataSink method), [1] (nvalchemi.dynamics.GPUBuffer method) (nvalchemi.dynamics.HostMemory method) (nvalchemi.dynamics.ZarrData method) write_empty_chunks (nvalchemi.data.datapipes.ZarrArrayConfig attribute) Z ZarrData (class in nvalchemi.dynamics) zero() (nvalchemi.data.Batch method) (nvalchemi.dynamics.DataSink method), [1] (nvalchemi.dynamics.GPUBuffer method) (nvalchemi.dynamics.HostMemory method) (nvalchemi.dynamics.ZarrData method) zero_forces (nvalchemi.dynamics.hooks.FreezeAtomsHook attribute), [1] zero_gradients() (in module nvalchemi.training)
