Index _ | A | B | C | D | E | F | G | H | I | J | K | L | M | N | O | P | Q | R | S | T | U | V | W | Z _ __init__() (AdaLNZeroResidualBlock method) (Attention method) (AutoregressiveWrapper method) (AutoResumeConfig method) (BlockAlignedGaussianModulationND method) (BlockDiagonalLearnableOmegaSIRENKernelND method) (BlockDiagonalMultiOmegaSIRENKernelND method) (CausalConv1D method) (CKConvND method) (ClassificationWrapper method) (DistributedDepthwiseConv1d method) (DistributedDepthwiseConv2d method) (DistributedDepthwiseConv3d method) (DropPath method) (ExperimentConfig method) (ExponentialModulationND method) (FiLMMonitorCallback method) (GaussianModulationND method) (GlobalResponseNorm method) (Hyena method) (IterationSpeedCallback method) (KernelFiLMGenerator method) (L2Norm method) (LayerScale method) (LazyConfig method) (LearnableOmegaSIRENKernelND method) (LearnableOmegaSIRENPositionalEmbeddingND method) (LightningWrapperBase method) (Mamba method) (MaskMonitorCallback method) (MLP method) (MultiOmegaSIRENKernelND method) (MultiOmegaSIRENPositionalEmbeddingND method) (OmegaScaleMonitorCallback method) (Patchify method) (PerHeadRMSNorm method) (PositionEmbeddingND method) (QKVConditionMixer method) (QKVSequenceMixer method) (RandomFourierKernelND method) (RandomFourierPositionalEmbeddingND method) (RegisterCompressConcat method) (RegisterPooling method) (RegressionWrapper method) (ResidualBlock method) (ResidualNetwork method) (RMSNorm method) (RMSNormChannelFirst method) (SchedulerConfig method) (Sequence1DVisualizationCallback method) (SIRENKernelND method) (SIRENPositionalEmbeddingND method) (StartFromCheckpointConfig method) (SubqOpsCausalConv1d method) (TrainConfig method) (TrainerConfig method) (UNetConvNext method) (UNetConvNextV2 method) (Unpatchify method) (ValidationImageGridCallback method) (ValidationVolumeGridCallback method) (ViT5Attention method) (ViT5ClassificationNet method) (ViT5HyenaAdapter method) (ViT5ResidualBlock method) (WalltimeCheckpointer method) (WandbCacheCleanupCallback method) (WandbConfig method) (WELLRegressionWrapper method) (WellUNetConvNext method) (WellUNetConvNextV2 method) _is_tuple_mode (CKConvND attribute) _periodic_per_axis (CKConvND attribute) _rope_ndim (Attention attribute) A accumulate_grad_steps (TrainConfig attribute) activation (MLP attribute) AdaLNZeroResidualBlock (class in nvsubquadratic.modules.residual_block) alias (AutoResumeConfig attribute) (StartFromCheckpointConfig attribute) AllToAllSingleFunction (class in nvsubquadratic.parallel.a2a_comms) apply_convolution() (CKConvND method) apply_qk_norm (Attention attribute) apply_qk_norm() (in module nvsubquadratic.utils.qk_norm) apply_rope_1d_bhl() (in module nvsubquadratic.utils.rope) apply_rope_1d_blh() (in module nvsubquadratic.utils.rope) apply_rope_2d_bhl() (in module nvsubquadratic.utils.rope) apply_rope_2d_blh() (in module nvsubquadratic.utils.rope) apply_rope_3d_bhl() (in module nvsubquadratic.utils.rope) apply_rope_3d_blh() (in module nvsubquadratic.utils.rope) Attention (class in nvsubquadratic.modules.attention) attn_dropout (Attention attribute) (ViT5Attention attribute) AutoregressiveWrapper (class in experiments.lightning_wrappers.autoregressive_wrapper) autoresume (ExperimentConfig attribute) AutoResumeConfig (class in experiments.default_cfg) B b2b_causal_conv1d() (in module nvsubquadratic.ops.causal_conv1d_custom) backend (MLP attribute) backward() (AllToAllSingleFunction static method) batch_size (TrainConfig attribute) beta (GlobalResponseNorm attribute) bias (DistributedDepthwiseConv1d attribute) (DistributedDepthwiseConv2d attribute) (DistributedDepthwiseConv3d attribute) bidirectional (Mamba attribute) binary_prediction() (ClassificationWrapper static method) BlockAlignedGaussianModulationND (class in nvsubquadratic.modules.masks_nd) BlockDiagonalLearnableOmegaSIRENKernelND (class in nvsubquadratic.modules.kernels_nd) BlockDiagonalMultiOmegaSIRENKernelND (class in nvsubquadratic.modules.kernels_nd) blocks (ResidualNetwork attribute) C callbacks (ExperimentConfig attribute) (StartFromCheckpointConfig attribute) causal (DistributedDepthwiseConv1d attribute) causal_conv1d() (in module nvsubquadratic.ops.causal_conv1d_custom) causal_fftconv1d_bhl() (in module nvsubquadratic.ops.fftconv_custom) causal_fftconv1d_bhl_w_reshape() (in module nvsubquadratic.ops.fftconv_custom) causal_fftconv1d_blh() (in module nvsubquadratic.ops.fftconv_custom) causal_fftconv1d_fp32_bhl() (in module nvsubquadratic.ops.fftconv) causal_fftconv1d_fp32_blh() (in module nvsubquadratic.ops.fftconv) CausalConv1D (class in nvsubquadratic.modules.causal_conv1d) channels_first (L2Norm property) (RMSNormChannelFirst attribute) check_val_every_n_epoch (TrainerConfig attribute) check_val_every_n_iterations (TrainerConfig attribute) checkpoint_every_n_steps (TrainerConfig attribute) checkpoint_monitor (TrainerConfig attribute) chunking_enabled() (in module nvsubquadratic.ops.fftconv_chunked) circular_fftconv1d_fp32_bhl() (in module nvsubquadratic.ops.circular_fftconv) circular_fftconv2d_fp32_bhl() (in module nvsubquadratic.ops.circular_fftconv) circular_fftconv3d_fp32_bhl() (in module nvsubquadratic.ops.circular_fftconv) CKConvND (class in nvsubquadratic.modules.ckconv_nd) ClassificationResNet (class in nvsubquadratic.networks.classification_resnet) ClassificationWrapper (class in experiments.lightning_wrappers.classification_wrapper) comment (ExperimentConfig attribute) compile (ExperimentConfig attribute) compile_compatible_fftconv (ExperimentConfig attribute) compile_mode (ExperimentConfig attribute) compress (RegisterCompressConcat attribute) compressed_dim (RegisterCompressConcat attribute) compute_relative_error() (in module nvsubquadratic.testing.utils) condition_in_proj (ResidualNetwork attribute) condition_mixer (ResidualBlock attribute) condition_mixer_norm (ResidualBlock attribute) condition_norm (AdaLNZeroResidualBlock attribute) condition_proj (AdaLNZeroResidualBlock attribute) configure_optimizers() (LightningWrapperBase method) construct_optimizer() (in module experiments.lightning_wrappers.base_lightning_wrapper) construct_rope_1d_cache_bhl() (in module nvsubquadratic.utils.rope) construct_rope_1d_cache_blh() (in module nvsubquadratic.utils.rope) construct_rope_2d_cache_bhl() (in module nvsubquadratic.utils.rope) construct_rope_2d_cache_blh() (in module nvsubquadratic.utils.rope) construct_rope_3d_cache_bhl() (in module nvsubquadratic.utils.rope) construct_rope_3d_cache_blh() (in module nvsubquadratic.utils.rope) construct_scheduler() (in module experiments.lightning_wrappers.base_lightning_wrapper) construct_trainer() (in module experiments.trainer) conv (Patchify attribute) core_layer (Mamba attribute) core_layer_rev (Mamba attribute) cuda_supports_quack() (in module nvsubquadratic.utils.quack_utils) D dali_imagenet_fused (in module experiments.datamodules) data_dim (CKConvND attribute) (Patchify attribute) (PositionEmbeddingND attribute) (RandomFourierKernelND attribute) (RandomFourierPositionalEmbeddingND attribute) (ResidualNetwork attribute) (SIRENKernelND attribute) (SIRENPositionalEmbeddingND attribute) (Unpatchify attribute) data_embeddings (PositionEmbeddingND attribute) dataset (ExperimentConfig attribute) debug (ExperimentConfig attribute) decay_iterations_percentage (SchedulerConfig attribute) deconv (Unpatchify attribute) deterministic (ExperimentConfig attribute) device (ExperimentConfig attribute) dim (GlobalResponseNorm attribute) (L2Norm attribute) distributed (LightningWrapperBase attribute) DistributedDepthwiseConv1d (class in nvsubquadratic.modules.distributed_depthwise_conv_nd) DistributedDepthwiseConv2d (class in nvsubquadratic.modules.distributed_depthwise_conv_nd) DistributedDepthwiseConv3d (class in nvsubquadratic.modules.distributed_depthwise_conv_nd) do (TrainConfig attribute) download_checkpoint() (in module experiments.utils.checkpointing) drop_path (ViT5ResidualBlock attribute) drop_path() (in module nvsubquadratic.modules.drop_path) drop_prob (DropPath attribute) dropout (AdaLNZeroResidualBlock attribute) (MLP attribute) (ResidualBlock attribute) dropout_in (ResidualNetwork attribute) DropPath (class in nvsubquadratic.modules.drop_path) E embedding_dim (PositionEmbeddingND attribute) (RandomFourierKernelND attribute) (RandomFourierPositionalEmbeddingND attribute) (SIRENKernelND attribute) (SIRENPositionalEmbeddingND attribute) emnist (in module experiments.datamodules) enable_chunking() (in module nvsubquadratic.ops.fftconv_chunked) enabled (AutoResumeConfig attribute) entity (WandbConfig attribute) eps (GlobalResponseNorm attribute) (L2Norm attribute) (RMSNorm attribute) (RMSNormChannelFirst attribute) eta_min (SchedulerConfig attribute) experiment_dir (ExperimentConfig attribute) ExperimentConfig (class in experiments.default_cfg) ExponentialModulationND (class in nvsubquadratic.modules.masks_nd) extra_repr() (Attention method) (CKConvND method) (DropPath method) (ExponentialModulationND method) (GaussianModulationND method) (Hyena method) (L2Norm method) (LearnableOmegaSIRENPositionalEmbeddingND method) (PerHeadRMSNorm method) (RegisterCompressConcat method) (RegisterPooling method) (ViT5Attention method) (ViT5HyenaAdapter method) F fft_backend (CKConvND attribute) fft_padding (CKConvND attribute) fftconv1d_fp32_bhl() (in module nvsubquadratic.ops.fftconv) fftconv1d_fp32_blh() (in module nvsubquadratic.ops.fftconv) fftconv2d_bhl() (in module nvsubquadratic.ops.fftconv_custom) fftconv2d_bhl_w_reshape() (in module nvsubquadratic.ops.fftconv_custom) fftconv2d_blh() (in module nvsubquadratic.ops.fftconv_custom) fftconv2d_fp32_bhl() (in module nvsubquadratic.ops.fftconv) fftconv2d_fp32_blh() (in module nvsubquadratic.ops.fftconv) fftconv3d_fp32_bhl() (in module nvsubquadratic.ops.fftconv) fftconv3d_fp32_blh() (in module nvsubquadratic.ops.fftconv) fftconv_fn (CKConvND attribute) fftconv_fn_bhl_input (CKConvND attribute) film_after_pos_embed (SIRENKernelND attribute) film_generator (SIRENKernelND attribute) FiLMMonitorCallback (class in experiments.callbacks.film_monitor) find_unused_parameters (TrainerConfig attribute) flop_count() (CKConvND method) (DropPath method) (GlobalResponseNorm method) (Hyena method) (KernelFiLMGenerator method) (LayerScale method) (MLP method) (PerHeadRMSNorm method) (QKVSequenceMixer method) (RegisterCompressConcat method) (RegisterPooling method) (RMSNorm method) (RMSNormChannelFirst method) (SIRENKernelND method) (ViT5Attention method) (ViT5ClassificationNet method) (ViT5HyenaAdapter method) (ViT5ResidualBlock method) forward() (AdaLNZeroResidualBlock method) (AllToAllSingleFunction static method) (Attention method) (CausalConv1D method) (CKConvND method) (ClassificationResNet method) (DistributedDepthwiseConv1d method) (DistributedDepthwiseConv2d method) (DistributedDepthwiseConv3d method) (DropPath method) (ExponentialModulationND method) (GaussianModulationND method) (GlobalResponseNorm method) (Hyena method) (KernelFiLMGenerator method) (L2Norm method) (LayerScale method) (LearnableOmegaSIRENPositionalEmbeddingND method) (LightningWrapperBase method) (Mamba method) (MLP method) (Patchify method) (PerHeadRMSNorm method) (PositionEmbeddingND method) (QKVConditionMixer method) (QKVSequenceMixer method) (RandomFourierKernelND method) (RandomFourierPositionalEmbeddingND method) (RegisterCompressConcat method) (RegisterPooling method) (ResidualBlock method) (ResidualNetwork method) (RMSNorm method) (RMSNormChannelFirst method) (Sine method) (SIRENKernelND method) (SIRENPositionalEmbeddingND method) (SoftTargetCrossEntropy method) (SubqOpsCausalConv1d method) (UNetConvNext method) (UNetConvNextV2 method) (Unpatchify method) (ViT5Attention method) (ViT5ClassificationNet method) (ViT5HyenaAdapter method) (ViT5ResidualBlock method) (WellUNetConvNext method) (WellUNetConvNextV2 method) G gamma (GlobalResponseNorm attribute) (LayerScale attribute) gate_nonlinear (Hyena attribute) gate_nonlinear_2 (Hyena attribute) GaussianModulationND (class in nvsubquadratic.modules.masks_nd) generate() (AutoregressiveWrapper method) generate_greedy() (AutoregressiveWrapper method) get_default_chunk_size() (in module nvsubquadratic.ops.fftconv_chunked) get_deterministic_run_name() (in module experiments.utils.cli) global_conv (Hyena attribute) GlobalResponseNorm (class in nvsubquadratic.modules.grn) grad_clip (TrainConfig attribute) gradient_checkpointing (ResidualNetwork attribute) grid_cache (RandomFourierPositionalEmbeddingND attribute) (SIRENPositionalEmbeddingND attribute) grid_type (CKConvND attribute) grid_w (ViT5HyenaAdapter attribute) grn (ViT5ResidualBlock attribute) group (AllToAllSingleFunction.ctx attribute) group_dim (DistributedDepthwiseConv1d attribute) (DistributedDepthwiseConv2d attribute) (DistributedDepthwiseConv3d attribute) H has_cls (ViT5Attention attribute) head_dim (Attention attribute) (PerHeadRMSNorm attribute) (ViT5Attention attribute) hidden_dim (Attention attribute) (CKConvND attribute) (DistributedDepthwiseConv1d attribute) (DistributedDepthwiseConv2d attribute) (DistributedDepthwiseConv3d attribute) (MLP attribute) (RegisterCompressConcat attribute) (ResidualNetwork attribute) (ViT5Attention attribute) hidden_linears (SIRENKernelND attribute) hidden_omega_0 (SIRENKernelND attribute) Hyena (class in nvsubquadratic.modules.hyena_nd) I in_channels (ResidualNetwork attribute) in_proj (ResidualNetwork attribute) init_parallel_state() (in module nvsubquadratic.parallel.utils) init_weights() (DistributedDepthwiseConv1d method) (DistributedDepthwiseConv2d method) (DistributedDepthwiseConv3d method) inner_mixer (ViT5HyenaAdapter attribute) input_norm (ResidualBlock attribute) (ViT5ResidualBlock attribute) instantiate() (in module nvsubquadratic.lazy_config) is_causal (Attention attribute) (CausalConv1D attribute) (CKConvND attribute) is_glu_variant (MLP attribute) iterations (TrainConfig attribute) IterationSpeedCallback (class in experiments.callbacks.iteration_speed) J job_group (WandbConfig attribute) K k_norm (Hyena attribute) (ViT5Attention attribute) kernel (CKConvND attribute) kernel_hidden_dim (KernelFiLMGenerator attribute) kernel_network (RandomFourierKernelND attribute) kernel_size (DistributedDepthwiseConv1d attribute) (DistributedDepthwiseConv2d attribute) (DistributedDepthwiseConv3d attribute) KernelFiLMGenerator (class in nvsubquadratic.modules.film) kv_proj (QKVConditionMixer attribute) L L2Norm (class in nvsubquadratic.utils.qk_norm) L_cache (RandomFourierKernelND attribute) (RandomFourierPositionalEmbeddingND attribute) (SIRENKernelND attribute) (SIRENPositionalEmbeddingND attribute) L_cache_per_axis (RandomFourierKernelND attribute) (RandomFourierPositionalEmbeddingND attribute) (SIRENKernelND attribute) (SIRENPositionalEmbeddingND attribute) LabeledEMAWeightAveraging (class in experiments.callbacks.model_ema) layer1 (MLP attribute) layer2 (MLP attribute) LayerScale (class in nvsubquadratic.modules.layer_scale) LazyConfig (class in nvsubquadratic.lazy_config) LearnableOmegaSIRENKernelND (class in nvsubquadratic.modules.kernels_nd) LearnableOmegaSIRENPositionalEmbeddingND (class in nvsubquadratic.modules.kernels_nd) lightning_wrapper_class (ExperimentConfig attribute) LightningWrapperBase (class in experiments.lightning_wrappers.base_lightning_wrapper) limit_test_batches (TrainerConfig attribute) limit_val_batches (TrainerConfig attribute) linear (LearnableOmegaSIRENPositionalEmbeddingND attribute) (MultiOmegaSIRENPositionalEmbeddingND attribute) (RandomFourierPositionalEmbeddingND attribute) (SIRENPositionalEmbeddingND attribute) load (StartFromCheckpointConfig attribute) load_checkpoint_state_dict() (in module experiments.utils.checkpointing) load_config_from_file() (in module experiments.utils.cli) load_state_dict() (MaskMonitorCallback method) (OmegaScaleMonitorCallback method) (ResumableSequentialLR method) logger (LightningWrapperBase property) logits (RegisterPooling attribute) ls_attn (ViT5ResidualBlock attribute) ls_mlp (ViT5ResidualBlock attribute) M main() (in module experiments.run) Mamba (class in nvsubquadratic.modules.mamba_nd) mask (CKConvND attribute) MaskMonitorCallback (class in experiments.callbacks.mask_monitor) max_dim_lengths (PositionEmbeddingND attribute) min_lr_ratio (SchedulerConfig attribute) mixed_fftconv1d_fp32_bhl() (in module nvsubquadratic.ops.mixed_fftconv) mixed_fftconv1d_fp32_bhl_chunked() (in module nvsubquadratic.ops.mixed_fftconv) mixed_fftconv1d_fp32_bhl_w_reshape() (in module nvsubquadratic.ops.mixed_fftconv) mixed_fftconv2d_fp32_bhl() (in module nvsubquadratic.ops.mixed_fftconv) mixed_fftconv2d_fp32_bhl_chunked() (in module nvsubquadratic.ops.mixed_fftconv) mixed_fftconv2d_fp32_bhl_w_reshape() (in module nvsubquadratic.ops.mixed_fftconv) mixed_fftconv3d_fp32_bhl() (in module nvsubquadratic.ops.mixed_fftconv) mixed_fftconv3d_fp32_bhl_chunked() (in module nvsubquadratic.ops.mixed_fftconv) mixed_fftconv3d_fp32_bhl_w_reshape() (in module nvsubquadratic.ops.mixed_fftconv) mixer (QKVConditionMixer attribute) (QKVSequenceMixer attribute) mlp (AdaLNZeroResidualBlock attribute) MLP (class in nvsubquadratic.modules.mlp) mlp (KernelFiLMGenerator attribute) (ResidualBlock attribute) (ViT5ResidualBlock attribute) mlp_hidden_dim (RandomFourierKernelND attribute) (SIRENKernelND attribute) mlp_norm (AdaLNZeroResidualBlock attribute) (ResidualBlock attribute) (ViT5ResidualBlock attribute) mnist (in module experiments.datamodules) mode (SchedulerConfig attribute) module nvsubquadratic monitor (SchedulerConfig attribute) mp_sharing_strategy (ExperimentConfig attribute) multiclass_prediction() (ClassificationWrapper static method) MultiOmegaSIRENKernelND (class in nvsubquadratic.modules.kernels_nd) MultiOmegaSIRENPositionalEmbeddingND (class in nvsubquadratic.modules.kernels_nd) N name (SchedulerConfig attribute) net (ExperimentConfig attribute) network (LightningWrapperBase attribute) norm (PerHeadRMSNorm attribute) num_blocks (BlockDiagonalLearnableOmegaSIRENKernelND attribute) (BlockDiagonalMultiOmegaSIRENKernelND attribute) (ResidualNetwork attribute) num_film_layers (KernelFiLMGenerator attribute) (SIRENKernelND attribute) num_groups (DistributedDepthwiseConv1d attribute) (DistributedDepthwiseConv2d attribute) (DistributedDepthwiseConv3d attribute) num_heads (Attention attribute) (PerHeadRMSNorm attribute) (ViT5Attention attribute) num_layers (RandomFourierKernelND attribute) (SIRENKernelND attribute) num_nodes (ExperimentConfig attribute) num_patches_h (ViT5Attention attribute) num_patches_w (ViT5Attention attribute) num_registers (RegisterCompressConcat attribute) (ViT5Attention attribute) (ViT5ResidualBlock attribute) nvsubquadratic module O off_block_scale (BlockDiagonalLearnableOmegaSIRENKernelND attribute) (BlockDiagonalMultiOmegaSIRENKernelND attribute) omega_0 (LearnableOmegaSIRENPositionalEmbeddingND attribute) (MultiOmegaSIRENKernelND attribute) (MultiOmegaSIRENPositionalEmbeddingND attribute) (RandomFourierKernelND attribute) (RandomFourierPositionalEmbeddingND attribute) (SIRENKernelND attribute) (SIRENPositionalEmbeddingND attribute) omega_0_const (LearnableOmegaSIRENPositionalEmbeddingND attribute) omega_0_per_block (BlockDiagonalLearnableOmegaSIRENKernelND attribute) (BlockDiagonalMultiOmegaSIRENKernelND attribute) omega_0_per_row (MultiOmegaSIRENKernelND attribute) (MultiOmegaSIRENPositionalEmbeddingND attribute) omega_0_scale (LearnableOmegaSIRENPositionalEmbeddingND attribute) omega_0_scale_max (LearnableOmegaSIRENPositionalEmbeddingND attribute) omega_0_scale_min (LearnableOmegaSIRENPositionalEmbeddingND attribute) OmegaScaleMonitorCallback (class in experiments.callbacks.omega_scale_monitor) on_after_backward() (IterationSpeedCallback method) (LightningWrapperBase method) on_before_backward() (IterationSpeedCallback method) (LightningWrapperBase method) on_before_optimizer_step() (LightningWrapperBase method) on_fit_end() (FiLMMonitorCallback method) on_fit_start() (FiLMMonitorCallback method) (LightningWrapperBase method) (MaskMonitorCallback method) (OmegaScaleMonitorCallback method) (WandbCacheCleanupCallback method) on_load_checkpoint() (ClassificationWrapper method) (LightningWrapperBase method) on_save_checkpoint() (ClassificationWrapper method) on_train_batch_end() (FiLMMonitorCallback method) (IterationSpeedCallback method) (MaskMonitorCallback method) (OmegaScaleMonitorCallback method) (Sequence1DVisualizationCallback method) (ValidationImageGridCallback method) (ValidationVolumeGridCallback method) (WalltimeCheckpointer method) on_train_batch_start() (IterationSpeedCallback method) on_train_epoch_end() (AutoregressiveWrapper method) (ClassificationWrapper method) (RegressionWrapper method) (WandbCacheCleanupCallback method) on_val_batch_end() (WalltimeCheckpointer method) on_validation_end() (IterationSpeedCallback method) on_validation_epoch_end() (AutoregressiveWrapper method) (ClassificationWrapper method) (LabeledEMAWeightAveraging method) (RegressionWrapper method) (Sequence1DVisualizationCallback method) (ValidationImageGridCallback method) (ValidationVolumeGridCallback method) (WELLRegressionWrapper method) on_validation_epoch_start() (LabeledEMAWeightAveraging method) on_validation_start() (IterationSpeedCallback method) optimizer (ExperimentConfig attribute) optimizer_cfg (LightningWrapperBase attribute) out_channels (ResidualNetwork attribute) out_dim (RandomFourierKernelND attribute) (RegisterCompressConcat property) (SIRENKernelND attribute) out_linear (RandomFourierKernelND attribute) (SIRENKernelND attribute) out_norm (ResidualNetwork attribute) out_proj (QKVConditionMixer attribute) (QKVSequenceMixer attribute) (ResidualNetwork attribute) output_norm (Hyena attribute) P parse_args() (in module experiments.run) partial_load (StartFromCheckpointConfig attribute) partial_wang_init_fn_with_num_layers() (in module nvsubquadratic.utils.init) patch_size (Patchify attribute) (Unpatchify attribute) Patchify (class in nvsubquadratic.modules.patchify) per_dim_embedding_dim (PositionEmbeddingND attribute) PerHeadRMSNorm (class in nvsubquadratic.modules.rms_norm) pixelhyena_norm (Hyena attribute) positional_embedding (BlockDiagonalLearnableOmegaSIRENKernelND attribute) (BlockDiagonalMultiOmegaSIRENKernelND attribute) (LearnableOmegaSIRENKernelND attribute) (MultiOmegaSIRENKernelND attribute) (RandomFourierKernelND attribute) (SIRENKernelND attribute) PositionEmbeddingND (class in nvsubquadratic.modules.position_encoding) precision (TrainConfig attribute) proj (ViT5Attention attribute) proj_drop (ViT5Attention attribute) project (WandbConfig attribute) Q q_norm (Hyena attribute) (ViT5Attention attribute) q_proj (QKVConditionMixer attribute) qk_norm (ViT5Attention attribute) qkv (ViT5Attention attribute) qkv_proj (QKVSequenceMixer attribute) QKVConditionMixer (class in nvsubquadratic.modules.condition_mixer) QKVSequenceMixer (class in nvsubquadratic.modules.sequence_mixer) R RandomFourierKernelND (class in nvsubquadratic.modules.kernels_nd) RandomFourierPositionalEmbeddingND (class in nvsubquadratic.modules.kernels_nd) reg_rope_base (ViT5Attention attribute) reg_rope_h (ViT5Attention attribute) reg_rope_w (ViT5Attention attribute) register_pooling (ViT5ResidualBlock attribute) register_start_idx (ViT5ResidualBlock attribute) RegisterCompressConcat (class in nvsubquadratic.modules.film) RegisterPooling (class in nvsubquadratic.modules.film) RegressionWrapper (class in experiments.lightning_wrappers.regression_wrapper) ResidualBlock (class in nvsubquadratic.modules.residual_block) ResidualNetwork (class in nvsubquadratic.networks.general_purpose_resnet) ResumableSequentialLR (class in nvsubquadratic.modules.schedulers) RMSNorm (class in nvsubquadratic.modules.rms_norm) RMSNormChannelFirst (class in nvsubquadratic.modules.rms_norm_channel_first) rope_base (Attention attribute) (ViT5Attention attribute) rope_cos (ViT5Attention attribute) rope_sin (ViT5Attention attribute) run_id (WandbConfig attribute) run_name (AutoResumeConfig attribute) run_path (StartFromCheckpointConfig attribute) run_start_time (TrainConfig attribute) run_time_limit_hours (TrainConfig attribute) S samples_per_epoch (TrainerConfig attribute) scale (Attention attribute) (ViT5Attention attribute) scheduler (ExperimentConfig attribute) scheduler_cfg (LightningWrapperBase attribute) SchedulerConfig (class in experiments.default_cfg) seed (ExperimentConfig attribute) Sequence1DVisualizationCallback (class in experiments.callbacks.sequence_visualization_1d) sequence_mixer (AdaLNZeroResidualBlock attribute) (ResidualBlock attribute) (ViT5ResidualBlock attribute) sequence_norm (AdaLNZeroResidualBlock attribute) set_default_chunk_size() (in module nvsubquadratic.ops.fftconv_chunked) setup_rank0_logging() (in module nvsubquadratic.parallel.utils) short_conv (Hyena attribute) shortcut (CKConvND attribute) Sine (class in nvsubquadratic.modules.kernels_nd) sine (SIRENKernelND attribute) SIRENKernelND (class in nvsubquadratic.modules.kernels_nd) SIRENPositionalEmbeddingND (class in nvsubquadratic.modules.kernels_nd) small_init() (in module nvsubquadratic.utils.init) SoftTargetCrossEntropy (class in experiments.lightning_wrappers.classification_wrapper) spatial_recall_dataset (in module experiments.datamodules) stable_iterations_percentage (SchedulerConfig attribute) start_from_checkpoint (ExperimentConfig attribute) StartFromCheckpointConfig (class in experiments.default_cfg) state_dict() (MaskMonitorCallback method) (OmegaScaleMonitorCallback method) state_key (MaskMonitorCallback property) (OmegaScaleMonitorCallback property) step_sizes (RandomFourierPositionalEmbeddingND attribute) (SIRENPositionalEmbeddingND attribute) strict (StartFromCheckpointConfig attribute) stride (Patchify attribute) (Unpatchify attribute) SubqOpsCausalConv1d (class in nvsubquadratic.modules.subq_ops_causal_conv1d) T tags (WandbConfig attribute) target_size (ResidualNetwork attribute) test_step() (AutoregressiveWrapper method) (ClassificationWrapper method) (RegressionWrapper method) (WELLRegressionWrapper method) tinyimagenet (in module experiments.datamodules) total_iterations (SchedulerConfig attribute) track_grad_norm (TrainConfig attribute) train (ExperimentConfig attribute) TrainConfig (class in experiments.default_cfg) trainer (ExperimentConfig attribute) TrainerConfig (class in experiments.default_cfg) training_step() (AutoregressiveWrapper method) (ClassificationWrapper method) (RegressionWrapper method) (WELLRegressionWrapper method) trunc_normal_init() (in module nvsubquadratic.utils.init) trunc_normal_init_factory() (in module nvsubquadratic.utils.init) type (AllToAllSingleFunction.ctx attribute) U UNetConvNext (class in nvsubquadratic.networks.baselines.unet_convnext) UNetConvNextV2 (class in nvsubquadratic.networks.baselines.unet_convnext_v2) Unpatchify (class in nvsubquadratic.modules.patchify) use_bias (RandomFourierPositionalEmbeddingND attribute) (SIRENPositionalEmbeddingND attribute) use_chunked_fftconv (CKConvND attribute) use_quack (RMSNorm attribute) (RMSNormChannelFirst attribute) use_rope (Attention attribute) V validation_step() (AutoregressiveWrapper method) (ClassificationWrapper method) (RegressionWrapper method) (WELLRegressionWrapper method) ValidationImageGridCallback (class in experiments.callbacks.image_grid_val_visualization) ValidationVolumeGridCallback (class in experiments.callbacks.image_grid_val_visualization) ViT5Attention (class in nvsubquadratic.modules.vit5_attention) ViT5ClassificationNet (class in nvsubquadratic.networks.vit5_classification) ViT5HyenaAdapter (class in nvsubquadratic.modules.vit5_hyena_adapter) ViT5ResidualBlock (class in nvsubquadratic.modules.vit5_residual_block) W WalltimeCheckpointer (class in experiments.callbacks.walltime_checkpointer) wandb (ExperimentConfig attribute) wandb_checkpoint_upload (TrainerConfig attribute) WandbCacheCleanupCallback (class in experiments.callbacks.wandb_cache_cleanup) WandbConfig (class in experiments.default_cfg) wang_init() (in module nvsubquadratic.utils.init) warmup_iterations_percentage (SchedulerConfig attribute) weight (DistributedDepthwiseConv1d attribute) (DistributedDepthwiseConv2d attribute) (DistributedDepthwiseConv3d attribute) (RMSNorm attribute) (RMSNormChannelFirst attribute) well (in module experiments.datamodules.pde) WELLRegressionWrapper (class in experiments.lightning_wrappers.well_lightning_wrapper) WellUNetConvNext (class in nvsubquadratic.networks.baselines.unet_convnext) WellUNetConvNextV2 (class in nvsubquadratic.networks.baselines.unet_convnext_v2) with_zigzag_splitting (AllToAllSingleFunction.ctx attribute) Z zigzag_gather_from_group_ranks() (in module nvsubquadratic.parallel.utils) zigzag_split_across_group_ranks() (in module nvsubquadratic.parallel.utils)
