Index

_

_act_fn_registry (in module xlstm_jax.models.xlstm_clean.components.feedforward) (in module xlstm_jax.models.xlstm_parallel.components.feedforward) (in module xlstm_jax.models.xlstm_pytorch.components.feedforward) _active_on_epochs (xlstm_jax.trainer.callbacks.callback.Callback attribute) _active_on_steps (xlstm_jax.trainer.callbacks.callback.Callback attribute) _batch_fn (xlstm_jax.dataset.grain_batch_rampup._BatchRampUpDatasetIterator attribute) (xlstm_jax.dataset.grain_batch_rampup.BatchRampUpIterDataset attribute) _batch_rampup_schedule (xlstm_jax.dataset.grain_batch_rampup._BatchRampUpDatasetIterator attribute) (xlstm_jax.dataset.grain_batch_rampup.BatchRampUpIterDataset attribute) _batch_step (xlstm_jax.dataset.grain_batch_rampup._BatchRampUpDatasetIterator attribute) _BatchRampUpDatasetIterator (class in xlstm_jax.dataset.grain_batch_rampup) _bias (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase property) _bias_ (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase attribute) _bias_ext2int() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) _bias_int2ext() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) _block_idx (xlstm_jax.models.xlstm_clean.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlockConfig attribute) _block_map (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) _build_global_shape_and_sharding() (in module xlstm_jax.dataset.multihost_dataloading) _check_for_nans() (xlstm_jax.trainer.logger.base_logger.Logger method) _check_input() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) _create_block_map() (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig method) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig method) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig method) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig method) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig method) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig method) _create_blocks() (xlstm_jax.models.xlstm_clean.xlstm_block_stack.BlockStack method) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.BlockStack method) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStack method) _create_weight_decay_optim_groups() (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1d method) (xlstm_jax.models.xlstm_pytorch.utils.WeightDecayOptimGroupMixin method) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel method) _dist_from_stddev() (in module xlstm_jax.models.shared.init) _document_borders (xlstm_jax.dataset.batch.LLMBatch attribute) (xlstm_jax.dataset.batch.LLMIndexedBatch attribute) _drop_remainder (xlstm_jax.dataset.grain_batch_rampup._BatchRampUpDatasetIterator attribute) (xlstm_jax.dataset.grain_batch_rampup.BatchRampUpIterDataset attribute) _dtype (xlstm_jax.models.llama.attention.SelfAttentionConfig property) (xlstm_jax.models.llama.feedforward.FeedForwardConfig property) (xlstm_jax.models.llama.llama.LlamaConfig property) (xlstm_jax.models.xlstm_clean.blocks.mlstm.cell.mLSTMCellConfig property) (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig property) (xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlockConfig property) (xlstm_jax.models.xlstm_clean.components.conv.CausalConv1dConfig property) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig property) (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpandConfig property) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig property) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig property) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig property) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig property) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig property) (xlstm_jax.models.xlstm_parallel.components.conv.CausalConv1dConfig property) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig property) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpandConfig property) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig property) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig property) _empty_val() (in module xlstm_jax.trainer.metrics) _eval_metric_shapes (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) _eval_model_in_train_loop() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) _every_n_epochs (xlstm_jax.trainer.callbacks.callback.Callback attribute) _every_n_steps (xlstm_jax.trainer.callbacks.callback.Callback attribute) _form_global_array() (in module xlstm_jax.dataset.multihost_dataloading) _gate_dtype (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig property) _get_annotations_full_dataclass() (in module xlstm_jax.configs) _get_chunkwise_fwbw_kernel() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3.triton_fwbw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice.triton_fwbw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize.triton_fwbw) (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef.triton_fwbw) _get_final_state() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell static method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda static method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla static method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase static method) _get_positions() (xlstm_jax.dataset.grain_transforms.InferSegmentations static method) _get_state() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) _get_weight_decay_optim_groups_for_modules() (xlstm_jax.models.xlstm_pytorch.utils.WeightDecayOptimGroupMixin static method) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel static method) _impl() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) _impl_step() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) _inner_embedding_dim (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) _instantiate_class_empty() (in module xlstm_jax.configs) _is_cuda() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method)

_jax_to_triton_dtype (in module xlstm_jax.kernels.kernel_utils) _lazy_init() (xlstm_jax.dataset.grain_transforms.HFTokenize method) _lm_head_dtype (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig property) _log_config() (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLogger method) _log_lr() (xlstm_jax.trainer.callbacks.lr_monitor.LearningRateMonitor method) _main_process_only (xlstm_jax.trainer.callbacks.callback.Callback attribute) _make_batch() (in module xlstm_jax.dataset.grain_batch_rampup) _mapped_data (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling attribute) _mlstm_backend_registry (in module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend) (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend) (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend) _mlstm_chunkwise__parallel_bw_dQKV() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_bw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_bw) (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_bw) _mlstm_chunkwise__parallel_fw_H() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_fw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_fw) (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_fw) _mlstm_chunkwise__recurrent_bw_dC() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_bw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_bw) (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_bw) _mlstm_chunkwise__recurrent_fw_C() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_fw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_fw) (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_fw) _mlstm_chunkwise_bw() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_bw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_bw) (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_bw) _mlstm_chunkwise_fw() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_fw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_fw) (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_fw) _mlstm_chunkwise_fwbw_generator() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3.triton_fwbw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice.triton_fwbw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize.triton_fwbw) (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef.triton_fwbw) _num_blocks (xlstm_jax.models.xlstm_clean.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) _out_features (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpandConfig attribute) _pad_array_to_shape() (in module xlstm_jax.dataset.multihost_dataloading) _pad_batch_multiple() (in module xlstm_jax.dataset.lmeval_pipeline) _parent (xlstm_jax.dataset.grain_batch_rampup._BatchRampUpDatasetIterator attribute) _parse_python_arguments_to_json() (in module xlstm_jax.configs) _parsed_string() (in module xlstm_jax.configs) _permute_input() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) _permute_output() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) _proj_up_dim (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_clean.utils.UpProjConfigMixin attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.utils.UpProjConfigMixin attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.utils.UpProjConfigMixin attribute) _python_dtype_to_cuda_dtype (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) _recurrent_kernel (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase property) _recurrent_kernel_ (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase attribute) _recurrent_kernel_ext2int() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) _recurrent_kernel_int2ext() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) _registry (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config_utils.NameAndKwargs attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.mLSTMBackendNameAndKwargs attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils.NameAndKwargs attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.mLSTMBackendNameAndKwargs attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config_utils.NameAndKwargs attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.mLSTMBackendNameAndKwargs attribute) _run_model_eval() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) _set_proj_up_dim() (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig method) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig method) (xlstm_jax.models.xlstm_clean.utils.UpProjConfigMixin method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig method) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig method) (xlstm_jax.models.xlstm_parallel.utils.UpProjConfigMixin method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig method) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig method) (xlstm_jax.models.xlstm_pytorch.utils.UpProjConfigMixin method) _step_counter (xlstm_jax.dataset.synthetic_dataloading.SyntheticDataIterator attribute) _tokenize() (xlstm_jax.dataset.grain_transforms.HFTokenize method) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling method) _train_metric_shapes (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) _update_single_metric() (in module xlstm_jax.trainer.metrics) _zero_state() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method)

A

accumulate_gradients() (in module xlstm_jax.distributed.single_gpu) accumulate_gradients_loop() (in module xlstm_jax.distributed.single_gpu) accumulate_gradients_scan() (in module xlstm_jax.distributed.single_gpu) act_fn (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.GatedFeedForward attribute) activation_function (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention.mLSTMBackendAttentionConfig attribute) add_bos (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) add_bos_token (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling attribute) add_embedding_dropout (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) add_eod (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) (xlstm_jax.dataset.grain_transforms.AddEODToken attribute) (xlstm_jax.dataset.grain_transforms.HFTokenize attribute) add_eos (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) add_eos_token (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling attribute) add_post_blocks_norm (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) add_post_norm (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute)

add_qk_norm (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) AddEODToken (class in xlstm_jax.dataset.grain_transforms) ademamix() (in module xlstm_jax.trainer.optimizer.ademamix) aggregate_metrics() (in module xlstm_jax.trainer.metrics) (lmeval_extended_evaluation.LMEvalEvaluation method) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) alpha (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] alpha_scheduler() (in module xlstm_jax.trainer.optimizer.ademamix) apply_chat_template (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) apply_fn (xlstm_jax.trainer.callbacks.extended_evaluation.EvalState attribute) apply_mlstm_param_reshapes() (in module xlstm_jax.utils.model_param_handling.handle_mlstm_simple) apply_rotary_emb() (in module xlstm_jax.models.llama.attention) (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention) apply_weight_transforms_() (in module xlstm_jax.utils.model_param_handling.convert_state_dict) assert_equal() (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_fw) assign_model_config_params() (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.simple.mLSTMBackendJaxConfig method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention.mLSTMBackendAttentionConfig method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbwConfig method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent.mLSTMBackendRecurrentConfig method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton.mLSTMBackendRecurrentTritonConfig method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple.mLSTMBackendParallelConfig method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTritonConfig method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple.mLSTMBackendTorchConfig method) async_gather() (in module xlstm_jax.distributed.tensor_parallel) async_gather_bidirectional() (in module xlstm_jax.distributed.tensor_parallel) async_gather_split() (in module xlstm_jax.distributed.tensor_parallel) async_scatter() (in module xlstm_jax.distributed.tensor_parallel) async_scatter_split() (in module xlstm_jax.distributed.tensor_parallel) attention() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention) attention_backend (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.attention.SelfAttentionConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) AttentionBackend (in module xlstm_jax.models.llama.attention) autocast_dtype (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTritonConfig attribute) axis (xlstm_jax.dataset.grain_transforms.ShiftData attribute)

B

backend (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) backend_fn (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell attribute) backend_fn_step (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell attribute) backend_name (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTritonConfig attribute) BackendNameType (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels) BackendType (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend) backward() (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils.CumsumFunction static method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils.ReversedCumsumFunction static method) base_dir (xlstm_jax.define_hydra_schemas.Config attribute) BaseLoggerConfig (class in xlstm_jax.define_hydra_schemas) Batch (class in xlstm_jax.dataset.batch) batch_class (xlstm_jax.dataset.grain_transforms.CollateToBatch attribute) batch_dataset_with_rampup() (in module xlstm_jax.dataset.grain_batch_rampup) batch_partition_specs (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) batch_rampup_factors (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) batch_size (lmeval_extended_evaluation.LMEvalEvaluation attribute) (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling attribute) (xlstm_jax.dataset.lmeval_pipeline.PadSequenceInBatchDataset attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) batch_size_per_device (xlstm_jax.define_hydra_schemas.Config attribute) batch_to_input() (xlstm_jax.trainer.base.trainer.TrainerModule static method) (xlstm_jax.trainer.llm.trainer.LLMTrainer static method) BatchRampUpIterDataset (class in xlstm_jax.dataset.grain_batch_rampup) benchmark_model() (in module xlstm_jax.models.xlstm_parallel.benchmark) beta1 (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] beta2 (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] beta3 (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] beta3_scheduler() (in module xlstm_jax.trainer.optimizer.ademamix) bias (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.ln.LayerNorm attribute) (xlstm_jax.models.xlstm_pytorch.components.ln.MultiHeadLayerNorm attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute)

bias_init (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpand attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpand attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) bias_linspace_init() (in module xlstm_jax.models.xlstm_clean.components.init) (in module xlstm_jax.models.xlstm_parallel.components.init) bias_linspace_init_() (in module xlstm_jax.models.xlstm_pytorch.components.init) block_map (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig property) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig property) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig property) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig property) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig property) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig property) blocks (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStack attribute) BlockStack (class in xlstm_jax.models.xlstm_clean.xlstm_block_stack) (class in xlstm_jax.models.xlstm_parallel.xlstm_block_stack) bootstrap_iters (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) bos_token_id (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling attribute) build_gradient_transformations() (in module xlstm_jax.trainer.optimizer.optimizer) build_lr_scheduler() (in module xlstm_jax.trainer.optimizer.scheduler) build_model() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) build_optimizer() (in module xlstm_jax.trainer.optimizer.optimizer) build_optimizer_function() (in module xlstm_jax.trainer.optimizer.optimizer) bytestring_to_sequence() (xlstm_jax.dataset.grain_transforms.ParseTokenizedArrayRecords static method)

C

cache_requests (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) Callback (class in xlstm_jax.trainer.callbacks.callback) CallbackConfig (class in xlstm_jax.trainer.callbacks.callback) callbacks (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) can_vmap_over_heads (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention.mLSTMBackendAttention property) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config.mLSTMBackend property) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbw property) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent.mLSTMBackendRecurrent property) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton.mLSTMBackendRecurrentTriton property) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple.mLSTMBackendParallel property) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTriton property) causal (xlstm_jax.models.llama.attention.SelfAttentionConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) causal_conv_bias (xlstm_jax.models.xlstm_clean.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_parallel.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1dConfig attribute) causal_forget_matrix() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw) (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw) CausalConv1d (class in xlstm_jax.models.xlstm_clean.components.conv) (class in xlstm_jax.models.xlstm_parallel.components.conv) (class in xlstm_jax.models.xlstm_pytorch.components.conv) CausalConv1dConfig (class in xlstm_jax.models.xlstm_clean.components.conv) (class in xlstm_jax.models.xlstm_parallel.components.conv) (class in xlstm_jax.models.xlstm_pytorch.components.conv) cell_norm_eps (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) cell_norm_type (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) cell_norm_type_v1 (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) channel_mixing (xlstm_jax.models.xlstm_clean.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_parallel.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1dConfig attribute) check_for_nan (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) check_val_every_n_epoch (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) check_val_every_n_steps (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) checkpoint() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) checkpoint_path (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint attribute) checkpointing (xlstm_jax.define_hydra_schemas.Config attribute) chunk_cumsum_bwd() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) chunk_cumsum_bwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) chunk_cumsum_fwd() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) chunk_cumsum_fwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) chunk_mlstm_bwd_kernel_dC() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.triton_chunk) chunk_mlstm_bwd_kernel_dqkvif() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.triton_chunk) chunk_mlstm_fwd_kernel_C() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.triton_chunk) chunk_mlstm_fwd_kernel_h() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.triton_chunk) chunk_reversed_cumsum_bwd() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) chunk_reversed_cumsum_bwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) chunk_reversed_cumsum_fwd() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) chunk_reversed_cumsum_fwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) chunk_size (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbwConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTritonConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) class_to_name() (in module xlstm_jax.import_utils) clip_by_global_norm_sharded() (in module xlstm_jax.trainer.optimizer.optimizer) clone() (xlstm_jax.models.xlstm_pytorch.components.util.ParameterProxy method) cmd_logging_name (xlstm_jax.trainer.logger.base_logger.LoggerConfig attribute), [1] CollateToBatch (class in xlstm_jax.dataset.grain_transforms) column_name (xlstm_jax.dataset.grain_data_processing.PaddedDataset attribute) (xlstm_jax.dataset.grain_transforms.HFNormalizeFeatures attribute) (xlstm_jax.dataset.grain_transforms.HFTokenize attribute) (xlstm_jax.dataset.grain_transforms.NormalizeFeatures attribute) (xlstm_jax.dataset.grain_transforms.ParseArrayRecords attribute) (xlstm_jax.dataset.grain_transforms.ParseTokenizedArrayRecords attribute) CombinedModelConfig (class in xlstm_jax.define_hydra_schemas) CompleteLLMIndexedBatch (class in xlstm_jax.dataset.lmeval_pipeline) compute_chunkwise_log_gates_vecB_vecA() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._chunkwise_gates) conditional_decorator() (in module xlstm_jax.models.xlstm_pytorch.components.util) Config (class in xlstm_jax.define_hydra_schemas) config (xlstm_jax.dataset.synthetic_dataloading.SyntheticDataIterator attribute) (xlstm_jax.models.llama.attention.SelfAttention attribute) (xlstm_jax.models.llama.feedforward.FeedForward attribute) (xlstm_jax.models.llama.llama.FFNBlock attribute) (xlstm_jax.models.llama.llama.LlamaTransformer attribute) (xlstm_jax.models.llama.llama.SelfAttentionBlock attribute) (xlstm_jax.models.llama.llama.TransformerBlock attribute) (xlstm_jax.models.llama.llama.TransformerBlockStack attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config.mLSTMBackend attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.cell.mLSTMCell attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayer attribute) (xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlock attribute) (xlstm_jax.models.xlstm_clean.components.conv.CausalConv1d attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.GatedFeedForward attribute) (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpand attribute) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.BlockStack attribute) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStack attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModel attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config.mLSTMBackend attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent.mLSTMBackendRecurrent attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTriton attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCell attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMInnerLayer attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayer attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer_v1.mLSTMLayerV1 attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlock attribute) (xlstm_jax.models.xlstm_parallel.components.conv.CausalConv1d attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForward attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.GatedFeedForward attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpand attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.BlockStack attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStack attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModel attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config.mLSTMBackend attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbw attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlock attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlock attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer attribute) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlock attribute) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1d attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.GatedFeedForward attribute) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpand attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStack attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel attribute) (xlstm_jax.trainer.callbacks.callback.Callback attribute) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation attribute) (xlstm_jax.trainer.logger.base_logger.Logger attribute) (xlstm_jax.trainer.logger.file_logger.FileLogger attribute) (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLogger attribute) (xlstm_jax.trainer.logger.wandb_logger.WandBLogger attribute)

config_class (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.simple.mLSTMBackendJax attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention.mLSTMBackendAttention attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbw attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent.mLSTMBackendRecurrent attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton.mLSTMBackendRecurrentTriton attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple.mLSTMBackendParallel attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTriton attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbw attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple.mLSTMBackendTorch attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlock attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlock attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer attribute) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlock attribute) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1d attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.GatedFeedForward attribute) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpand attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStack attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel attribute) config_format (xlstm_jax.trainer.logger.file_logger.FileLoggerConfig attribute), [1] config_to_log (xlstm_jax.trainer.logger.file_logger.FileLogger attribute) (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLogger attribute) (xlstm_jax.trainer.logger.wandb_logger.WandBLogger attribute) ConfigDict (class in xlstm_jax.configs) constant_rampup_schedule() (in module xlstm_jax.dataset.grain_batch_rampup) constants (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) context_length (lmeval_extended_evaluation.LMEvalEvaluation attribute) (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.define_hydra_schemas.Config attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.simple.mLSTMBackendJaxConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention.mLSTMBackendAttentionConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent.mLSTMBackendRecurrentConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple.mLSTMBackendParallelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple.mLSTMBackendTorchConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) contiguous() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) conv1d (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) conv1d_kernel_size (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayerConfig attribute) conv1d_kwargs (xlstm_jax.models.xlstm_clean.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_parallel.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1dConfig attribute) conv1d_step() (in module xlstm_jax.models.xlstm_pytorch.components.conv) conv_act_fn (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) convert_mlstm_checkpoint_jax_to_torch_simple() (in module xlstm_jax.utils.model_param_handling.handle_mlstm_simple) convert_orbax_checkpoint_to_torch_state_dict() (in module xlstm_jax.utils.model_param_handling.convert_checkpoint) convert_state_dict_keys_() (in module xlstm_jax.utils.model_param_handling.convert_state_dict) cooldown_lr (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig attribute), [1] cooldown_steps (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig attribute), [1] count (xlstm_jax.trainer.optimizer.ademamix.ScaleByAdemamixState attribute) count_m2 (xlstm_jax.trainer.optimizer.ademamix.ScaleByAdemamixState attribute) create() (lmeval_extended_evaluation.LMEvalEvaluationConfig method) (xlstm_jax.trainer.callbacks.callback.CallbackConfig method) (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpointConfig method) (xlstm_jax.trainer.callbacks.extended_evaluation.EvalState class method) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluationConfig method) (xlstm_jax.trainer.callbacks.lr_monitor.LearningRateMonitorConfig method) (xlstm_jax.trainer.callbacks.profiler.JaxProfilerConfig method) (xlstm_jax.trainer.logger.base_logger.LoggerToolsConfig method) (xlstm_jax.trainer.logger.file_logger.FileLoggerConfig method) (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLoggerConfig method) (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig method) create_batch() (in module xlstm_jax.models.xlstm_parallel.benchmark) create_batch_rampup_schedule() (in module xlstm_jax.dataset.grain_batch_rampup) create_common_init_fn() (in module xlstm_jax.models.shared.init) create_data_iterator() (in module xlstm_jax.dataset.input_pipeline_interface) create_evaluation_step_function() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) create_feedforward() (in module xlstm_jax.models.xlstm_clean.components.feedforward) (in module xlstm_jax.models.xlstm_parallel.components.feedforward) (in module xlstm_jax.models.xlstm_pytorch.components.feedforward) create_full_state_dict_key_mapping() (in module xlstm_jax.utils.model_param_handling.convert_state_dict) create_jitted_functions() (lmeval_extended_evaluation.LMEvalEvaluation method) (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) create_layer() (in module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.layer_factory) (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.layer_factory) (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.layer_factory) create_mixed_data_iterator() (in module xlstm_jax.dataset.input_pipeline_interface) create_mlstm_backend (in module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend) (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend) (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend) create_mlstm_simple_config_from_jax_config() (in module xlstm_jax.utils.model_param_handling.handle_mlstm_simple) create_modified_exemplary_batch() (lmeval_extended_evaluation.LMEvalEvaluation method) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) create_name_instance() (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) create_recurrent_evaluation_step_function() (lmeval_extended_evaluation.LMEvalEvaluation method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) create_tokenizer_fn (xlstm_jax.dataset.grain_transforms.HFTokenize attribute) create_train_eval_configs() (xlstm_jax.dataset.configs.DataConfig class method) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig class method) (xlstm_jax.dataset.configs.HFHubDataConfig class method) (xlstm_jax.dataset.configs.SyntheticDataConfig class method) create_training_step_function() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) cumsum() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) CumsumFunction (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) curdir (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.cuda_init)

D

data_axis_name (xlstm_jax.models.configs.ParallelConfig attribute) data_axis_size (xlstm_jax.models.configs.ParallelConfig attribute) data_column (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) (xlstm_jax.dataset.grain_transforms.ParseFeatures attribute) data_config_type (xlstm_jax.dataset.configs.DataConfig attribute) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) (xlstm_jax.dataset.configs.SyntheticDataConfig attribute) data_eval (xlstm_jax.define_hydra_schemas.Config attribute) data_generator (xlstm_jax.dataset.synthetic_dataloading.SyntheticDataIterator attribute) data_module (xlstm_jax.trainer.callbacks.callback.Callback attribute) data_path (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) data_shuffle_seed (xlstm_jax.dataset.configs.DataConfig attribute) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) (xlstm_jax.dataset.configs.SyntheticDataConfig attribute) data_train (xlstm_jax.define_hydra_schemas.Config attribute) dataclass_kwonly (in module xlstm_jax.dataset.batch) DataConfig (class in xlstm_jax.dataset.configs) DataEvalConfig (class in xlstm_jax.define_hydra_schemas) DataIterator (in module xlstm_jax.dataset.input_pipeline_interface) (in module xlstm_jax.trainer.data_module) dataloader (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator attribute) dataloader_host_count (xlstm_jax.dataset.lmeval_pipeline.MultihostSortedRemapDataset attribute) dataloader_path (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint attribute) DataloaderModule (class in xlstm_jax.trainer.data_module) dataset (xlstm_jax.dataset.grain_data_processing.PaddedDataset attribute) (xlstm_jax.dataset.lmeval_pipeline.MultihostSortedRemapDataset attribute) (xlstm_jax.dataset.lmeval_pipeline.PadBatchDataset attribute) (xlstm_jax.dataset.lmeval_pipeline.PadSequenceInBatchDataset attribute) (xlstm_jax.dataset.lmeval_pipeline.SortedDataset attribute) (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) dataset_size (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator attribute) (xlstm_jax.dataset.synthetic_dataloading.SyntheticDataIterator property) DataTrainConfig (class in xlstm_jax.define_hydra_schemas) debug (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) debug_cell (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) decay_steps (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig attribute), [1] default_train_log_modes (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) defines (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig property) defines_to_cflags() (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.cuda_init) delete_arrays_in_pytree() (in module xlstm_jax.utils.pytree_utils) dense_fn (xlstm_jax.distributed.tensor_parallel.TPAsyncDense attribute), [1] (xlstm_jax.distributed.tensor_parallel.TPDense attribute), [1] dense_name (xlstm_jax.distributed.tensor_parallel.TPAsyncDense attribute), [1] (xlstm_jax.distributed.tensor_parallel.TPDense attribute), [1] device (xlstm_jax.define_hydra_schemas.Config attribute) device_count (xlstm_jax.define_hydra_schemas.Config attribute) device_metrics_aggregation() (in module xlstm_jax.trainer.callbacks.extended_evaluation) device_type (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) document_idx (xlstm_jax.dataset.batch.LLMIndexedBatch attribute) donate_train_state (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) drop_remainder (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) dropout (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.GatedFeedForward attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute)

dropout_rate (xlstm_jax.models.llama.attention.SelfAttentionConfig attribute) (xlstm_jax.models.llama.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) ds1 (xlstm_jax.define_hydra_schemas.DataEvalConfig attribute) (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) ds10 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) ds2 (xlstm_jax.define_hydra_schemas.DataEvalConfig attribute) (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) ds3 (xlstm_jax.define_hydra_schemas.DataEvalConfig attribute) (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) ds4 (xlstm_jax.define_hydra_schemas.DataEvalConfig attribute) (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) ds5 (xlstm_jax.define_hydra_schemas.DataEvalConfig attribute) (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) ds6 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) ds7 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) ds8 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) ds9 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) dtype (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.attention.SelfAttentionConfig attribute) (xlstm_jax.models.llama.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_clean.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) dtype_a (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) dtype_b (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) DTYPE_DICT (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) dtype_g (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) dtype_r (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) dtype_s (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) dtype_w (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) DTYPES (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell)

E

emb_dropout (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel attribute) embedding_dim (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) empty_llm_indexed_sample() (in module xlstm_jax.dataset.lmeval_pipeline) empty_sequence (xlstm_jax.dataset.grain_data_processing.PaddedDataset property) enable_async_checkpointing (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpointConfig attribute), [1] enable_automatic_mixed_precision (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) enable_progress_bar (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) end_epoch() (xlstm_jax.trainer.logger.base_logger.Logger method) end_lr (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig attribute), [1] end_lr_factor (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig attribute), [1] eod_token_id (xlstm_jax.dataset.grain_transforms.AddEODToken attribute) (xlstm_jax.dataset.grain_transforms.HFTokenize attribute) (xlstm_jax.dataset.grain_transforms.InferSegmentations attribute) (xlstm_jax.dataset.grain_transforms.ShiftData attribute) eos_token_id (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling attribute) epoch (xlstm_jax.trainer.logger.base_logger.Logger attribute)

epoch_start_time_stack (xlstm_jax.trainer.logger.base_logger.Logger attribute) eps (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent.mLSTMBackendRecurrentConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton.mLSTMBackendRecurrentTritonConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTritonConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.ln.LayerNorm attribute) (xlstm_jax.models.xlstm_pytorch.components.ln.MultiHeadLayerNorm attribute) (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] err_msg (in module xlstm_jax.kernels) eval_function() (lmeval_extended_evaluation.LMEvalEvaluation method) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) eval_model() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) eval_step (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation attribute) EvalState (class in xlstm_jax.trainer.callbacks.extended_evaluation) evaluation_tasks (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) every_n_epochs (xlstm_jax.trainer.callbacks.callback.CallbackConfig attribute) (xlstm_jax.trainer.callbacks.lr_monitor.LearningRateMonitorConfig attribute) (xlstm_jax.trainer.callbacks.profiler.JaxProfilerConfig attribute), [1] every_n_steps (xlstm_jax.trainer.callbacks.callback.CallbackConfig attribute) (xlstm_jax.trainer.callbacks.lr_monitor.LearningRateMonitorConfig attribute) (xlstm_jax.trainer.callbacks.profiler.JaxProfilerConfig attribute), [1] execute_pipeline() (in module xlstm_jax.distributed.pipeline_parallel) execute_pipeline_step() (in module xlstm_jax.distributed.pipeline_parallel) exmp_batch (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) expand_factor_up (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpandConfig attribute) ExtendedEvaluation (class in xlstm_jax.trainer.callbacks.extended_evaluation) ExtendedEvaluationConfig (class in xlstm_jax.trainer.callbacks.extended_evaluation) external_to_internal (xlstm_jax.models.xlstm_pytorch.components.util.ParameterProxy attribute) extra_repr() (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpand method) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpand method) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpand method)

F

feature_dim (xlstm_jax.models.xlstm_clean.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_parallel.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1dConfig attribute) FeedForward (class in xlstm_jax.models.llama.feedforward) (class in xlstm_jax.models.xlstm_parallel.components.feedforward) feedforward (xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlockConfig attribute) FeedForwardConfig (class in xlstm_jax.models.llama.feedforward) (class in xlstm_jax.models.xlstm_clean.components.feedforward) (class in xlstm_jax.models.xlstm_parallel.components.feedforward) (class in xlstm_jax.models.xlstm_pytorch.components.feedforward) ff_dtype (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) ff_type (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) ffn_dim_multiplier (xlstm_jax.models.llama.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) ffn_multiple_of (xlstm_jax.models.llama.llama.LlamaConfig attribute) FFNBlock (class in xlstm_jax.models.llama.llama) fgate (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer attribute) fgate_bias_init_range (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) file_logger_config_format (xlstm_jax.define_hydra_schemas.BaseLoggerConfig attribute) file_logger_log_dir (xlstm_jax.define_hydra_schemas.BaseLoggerConfig attribute) FileLogger (class in xlstm_jax.trainer.logger.file_logger) FileLoggerConfig (class in xlstm_jax.trainer.logger.file_logger) finalize() (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint method) (xlstm_jax.trainer.logger.base_logger.Logger method) (xlstm_jax.trainer.logger.base_logger.LoggerTool method) (xlstm_jax.trainer.logger.file_logger.FileLogger method) (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLogger method) (xlstm_jax.trainer.logger.wandb_logger.WandBLogger method) finalize_metrics() (lmeval_extended_evaluation.LMEvalEvaluation method) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) find_parameter_match_key() (in module xlstm_jax.utils.model_param_handling.convert_state_dict) first_step (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) flatten_dict() (in module xlstm_jax.models.xlstm_parallel.training) (in module xlstm_jax.utils.pytree_utils) flatten_pytree() (in module xlstm_jax.utils.pytree_utils) fold_rng_over_axis() (in module xlstm_jax.distributed.array_utils) forward() (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config.mLSTMBackend method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbw method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple.mLSTMBackendTorch method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils.CumsumFunction static method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils.ReversedCumsumFunction static method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlock method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlock method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer method) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlock method) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1d method) (xlstm_jax.models.xlstm_pytorch.components.feedforward.GatedFeedForward method) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpand method) (xlstm_jax.models.xlstm_pytorch.components.ln.LayerNorm method) (xlstm_jax.models.xlstm_pytorch.components.ln.MultiHeadLayerNorm method) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStack method) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel method)

forward_clipval (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) found_nans (xlstm_jax.trainer.logger.base_logger.Logger attribute) from_dict() (xlstm_jax.configs.ConfigDict static method) (xlstm_jax.models.configs.ModelConfig static method) (xlstm_jax.trainer.base.trainer.TrainerConfig static method) (xlstm_jax.trainer.callbacks.callback.CallbackConfig static method) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig static method) (xlstm_jax.trainer.logger.base_logger.LoggerConfig static method) (xlstm_jax.trainer.logger.base_logger.LoggerToolsConfig static method) (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig static method) (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig static method) from_inputs() (xlstm_jax.dataset.batch.LLMBatch static method) (xlstm_jax.dataset.batch.LLMIndexedBatch static method) from_metadata() (xlstm_jax.models.configs.ModelConfig static method) from_train_state() (xlstm_jax.trainer.callbacks.extended_evaluation.EvalState class method) fsdp_axis_name (xlstm_jax.models.configs.ParallelConfig attribute) fsdp_axis_size (xlstm_jax.models.configs.ParallelConfig attribute) fsdp_gather_dtype (xlstm_jax.models.configs.ParallelConfig attribute) fsdp_grad_scatter_dtype (xlstm_jax.models.configs.ParallelConfig attribute) fsdp_min_weight_size (xlstm_jax.models.configs.ParallelConfig attribute) fsdp_modules (xlstm_jax.models.configs.ParallelConfig attribute) full_dataset_length (xlstm_jax.dataset.grain_data_processing.PaddedDataset attribute) func (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbw attribute) function (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) fuse_document_results() (in module lmeval_extended_evaluation) fwbw_backward() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw) fwbw_forward() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw)

G

gate_dtype (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) gate_input (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) gate_linear_headwise (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) gate_soft_cap (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) GatedFeedForward (class in xlstm_jax.models.xlstm_clean.components.feedforward) (class in xlstm_jax.models.xlstm_parallel.components.feedforward) (class in xlstm_jax.models.xlstm_pytorch.components.feedforward) gather_array_with_mean_grads() (in module xlstm_jax.distributed.data_parallel) gather_params() (in module xlstm_jax.distributed.data_parallel) generate_tokens() (in module xlstm_jax.trainer.llm.sampling) get() (xlstm_jax.configs.ConfigDict method) (xlstm_jax.models.configs.ModelConfig method) (xlstm_jax.trainer.base.trainer.TrainerConfig method) (xlstm_jax.trainer.callbacks.callback.CallbackConfig method) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig method) (xlstm_jax.trainer.logger.base_logger.LoggerConfig method) (xlstm_jax.trainer.logger.base_logger.LoggerToolsConfig method) (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig method) (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig method) get_act_fn() (in module xlstm_jax.models.xlstm_clean.components.feedforward) (in module xlstm_jax.models.xlstm_parallel.components.feedforward) (in module xlstm_jax.models.xlstm_pytorch.components.feedforward) get_cache_shape_dtype_struct() (xlstm_jax.trainer.llm.trainer.LLMTrainer method) get_class_for_name() (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) get_config_class_for_kwargs() (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) get_document_borders() (xlstm_jax.dataset.batch.LLMBatch method) (xlstm_jax.dataset.batch.LLMIndexedBatch method) get_dtype_struct() (xlstm_jax.dataset.batch.LLMBatch static method) (xlstm_jax.dataset.batch.LLMIndexedBatch static method) get_generate_fn() (xlstm_jax.trainer.llm.trainer.LLMTrainer method) get_grad_norms() (in module xlstm_jax.trainer.base.param_utils) get_loglevel() (in module xlstm_jax.trainer.logger.cmd_logging) get_metric_postprocess_fn() (lmeval_extended_evaluation.LMEvalEvaluation static method) (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) get_metrics() (in module xlstm_jax.trainer.metrics) get_model_rng() (xlstm_jax.trainer.base.trainer.TrainerModule static method) (xlstm_jax.trainer.llm.trainer.LLMTrainer static method) get_next_batch_sharded() (in module xlstm_jax.dataset.multihost_dataloading)

get_num_params() (in module xlstm_jax.models.xlstm_parallel.training) (in module xlstm_jax.trainer.base.param_utils) (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) get_param_mask_fn() (in module xlstm_jax.trainer.base.param_utils) get_param_norms() (in module xlstm_jax.trainer.base.param_utils) get_partial_mLSTMBlock() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.block) get_process_loading_data() (in module xlstm_jax.dataset.input_pipeline_interface) get_sample() (xlstm_jax.dataset.batch.LLMBatch class method) (xlstm_jax.dataset.batch.LLMIndexedBatch class method) get_shape_dtype_pytree() (in module xlstm_jax.utils.pytree_utils) get_sharded_global_norm() (in module xlstm_jax.trainer.base.param_utils) get_sharded_norm_logits() (in module xlstm_jax.trainer.base.param_utils) get_state() (xlstm_jax.dataset.grain_batch_rampup._BatchRampUpDatasetIterator method) (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator method) get_stride() (in module xlstm_jax.kernels.stride_utils) get_strides() (in module xlstm_jax.kernels.stride_utils) get_tokenizer_vocab_size() (in module xlstm_jax.train_init_fns) get_train_step_fn() (in module xlstm_jax.models.xlstm_parallel.training) get_weight_decay_optim_group_param_names() (xlstm_jax.models.xlstm_pytorch.utils.WeightDecayOptimGroupMixin method) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel method) get_weight_decay_optim_groups() (xlstm_jax.models.xlstm_pytorch.utils.WeightDecayOptimGroupMixin method) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel method) global_batch_size (xlstm_jax.dataset.configs.DataConfig attribute) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) (xlstm_jax.dataset.configs.SyntheticDataConfig attribute) (xlstm_jax.dataset.lmeval_pipeline.MultihostSortedRemapDataset attribute) (xlstm_jax.define_hydra_schemas.Config attribute) global_mesh (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator attribute) global_step (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) grad (xlstm_jax.models.xlstm_pytorch.components.util.ParameterProxy property) grad_clip_norm (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] grad_clip_value (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] gradient_accumulate_scan (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) gradient_accumulate_steps (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) gradient_recurrent_clipval (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) gradient_recurrent_cut (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) GRAIN_AVAILABLE (in module xlstm_jax.dataset.input_pipeline_interface) grain_packing (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) grain_packing_bin_count (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) GrainArrayRecordsDataConfig (class in xlstm_jax.dataset.configs) greedy_sampling() (in module xlstm_jax.trainer.llm.sampling) group_norm (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer attribute) group_norm_weight (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayerConfig attribute) groups (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1d attribute)

H

head_dim (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.attention.SelfAttentionConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase property) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig property) hf_access_token (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) hf_cache_dir (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute)

hf_data_dir (xlstm_jax.dataset.configs.HFHubDataConfig attribute) hf_data_files (xlstm_jax.dataset.configs.HFHubDataConfig attribute) hf_num_data_processes (xlstm_jax.dataset.configs.HFHubDataConfig attribute) hf_path (xlstm_jax.dataset.configs.HFHubDataConfig attribute) HFHubDataConfig (class in xlstm_jax.dataset.configs) HFNormalizeFeatures (class in xlstm_jax.dataset.grain_transforms) HFPrefixTokenize (class in xlstm_jax.dataset.grain_transforms) HFTokenize (class in xlstm_jax.dataset.grain_transforms) HFTokenizeLogLikelihood (class in xlstm_jax.dataset.lmeval_dataset) HFTokenizeLogLikelihoodRolling (class in xlstm_jax.dataset.lmeval_dataset) hidden_size (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) HostMetricElement (in module xlstm_jax.common_types) HostMetrics (in module xlstm_jax.common_types)

I

igate (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer attribute) igate_bias_init_range (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) ImmutableMetricElement (in module xlstm_jax.common_types) ImmutableMetrics (in module xlstm_jax.common_types) in_features (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpandConfig attribute) InferSegmentations (class in xlstm_jax.dataset.grain_transforms) infinite_eval_chunksize (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) init_callbacks() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) init_data_iterator() (in module xlstm_jax.train_init_fns) init_distribution (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) init_distribution_embed (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) init_distribution_out (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) init_eval_metrics() (lmeval_extended_evaluation.LMEvalEvaluation method) (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) init_logger() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) init_logger_config() (in module xlstm_jax.train_init_fns) init_lr_monitor_config() (in module xlstm_jax.train_init_fns) init_mesh() (in module xlstm_jax.models.xlstm_parallel.benchmark) (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) init_mixed_data_iterator() (in module xlstm_jax.train_init_fns) init_model() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method)

init_model_checkpointing() (in module xlstm_jax.train_init_fns) init_model_config() (in module xlstm_jax.train_init_fns) init_optimizer() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) init_optimizer_config() (in module xlstm_jax.train_init_fns) init_parallel() (in module xlstm_jax.train_init_fns) init_profiler_config() (in module xlstm_jax.train_init_fns) init_scheduler_config() (in module xlstm_jax.train_init_fns) init_single_data_iterator() (in module xlstm_jax.train_init_fns) init_train_metrics() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) init_trainer() (in module xlstm_jax.train_init_fns) init_xlstm() (in module xlstm_jax.models.xlstm_parallel.training) InitDistribution (in module xlstm_jax.models.shared.init) InitFnName (in module xlstm_jax.models.shared.init) initial_val (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) initialize_mesh() (in module xlstm_jax.distributed.mesh_utils) input_dim (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig property) input_shape (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) inputs (xlstm_jax.dataset.batch.Batch attribute) (xlstm_jax.dataset.batch.LLMBatch attribute) (xlstm_jax.dataset.batch.LLMIndexedBatch attribute) inputs_position (xlstm_jax.dataset.batch.LLMBatch attribute) (xlstm_jax.dataset.batch.LLMIndexedBatch attribute) inputs_segmentation (xlstm_jax.dataset.batch.LLMBatch attribute) (xlstm_jax.dataset.batch.LLMIndexedBatch attribute) instance() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellCUDA class method) intermediates_log_modes (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) internal_input_shape (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) internal_to_external (xlstm_jax.models.xlstm_pytorch.components.util.ParameterProxy attribute) is_partitioned() (in module xlstm_jax.trainer.base.param_utils) iterator_length (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator attribute)

J

jax2triton_dtype() (in module xlstm_jax.kernels.kernel_utils)

JaxProfiler (class in xlstm_jax.trainer.callbacks.profiler) JaxProfilerConfig (class in xlstm_jax.trainer.callbacks.profiler)

K

k_proj (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) keep_c (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbwConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) keep_G (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbwConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) keep_gates (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbwConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) keep_M (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbwConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) kernel_init (xlstm_jax.distributed.tensor_parallel.TPAsyncDense attribute), [1] (xlstm_jax.distributed.tensor_parallel.TPDense attribute), [1] (xlstm_jax.models.shared.lm_head.TPLMHead attribute) (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpand attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpand attribute)

kernel_init_adjustment (xlstm_jax.distributed.tensor_parallel.TPAsyncDense attribute), [1] (xlstm_jax.distributed.tensor_parallel.TPDense attribute), [1] kernel_size (xlstm_jax.models.xlstm_clean.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_parallel.components.conv.CausalConv1dConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1dConfig attribute) key (xlstm_jax.dataset.lmeval_pipeline.SortedDataset attribute) key_map (xlstm_jax.dataset.grain_transforms.CollateToBatch attribute) kwargs (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config_utils.NameAndKwargs attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.mLSTMBackendNameAndKwargs attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils.NameAndKwargs attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.mLSTMBackendNameAndKwargs attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config_utils.NameAndKwargs attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.mLSTMBackendNameAndKwargs attribute)

L

last_step (xlstm_jax.trainer.logger.base_logger.Logger attribute) last_step_time (xlstm_jax.trainer.logger.base_logger.Logger attribute) layer_type (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) LayerNorm (class in xlstm_jax.models.xlstm_pytorch.components.ln) LayerNorm() (in module xlstm_jax.models.xlstm_clean.components.ln) learnable_skip (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) LearningRateMonitor (class in xlstm_jax.trainer.callbacks.lr_monitor) LearningRateMonitorConfig (class in xlstm_jax.trainer.callbacks.lr_monitor) limit_requests (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) LinearHeadwiseExpand (class in xlstm_jax.models.xlstm_clean.components.linear_headwise) (class in xlstm_jax.models.xlstm_parallel.components.linear_headwise) (class in xlstm_jax.models.xlstm_pytorch.components.linear_headwise) LinearHeadwiseExpandConfig (class in xlstm_jax.models.xlstm_clean.components.linear_headwise) (class in xlstm_jax.models.xlstm_parallel.components.linear_headwise) (class in xlstm_jax.models.xlstm_pytorch.components.linear_headwise) LlamaConfig (class in xlstm_jax.models.llama.llama) LlamaTransformer (class in xlstm_jax.models.llama.llama) LLMBatch (class in xlstm_jax.dataset.batch) LLMIndexedBatch (class in xlstm_jax.dataset.batch) LLMTrainer (class in xlstm_jax.trainer.llm.trainer) LLMTrainerConfig (class in xlstm_jax.trainer.llm.trainer) lm (lmeval_extended_evaluation.LMEvalEvaluation attribute) lm_head (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel attribute) lm_head_dtype (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.shared.lm_head.TPLMHead attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) lmeval_extended_evaluation module lmeval_preprocessing_pipeline() (in module xlstm_jax.dataset.lmeval_pipeline) LMEvalEvaluation (class in lmeval_extended_evaluation) LMEvalEvaluationConfig (class in lmeval_extended_evaluation) load() (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.cuda_init) load_array_record_dataset() (in module xlstm_jax.dataset.grain_data_processing) load_data_loaders() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) load_dataloader() (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint method) load_from_checkpoint() (xlstm_jax.trainer.base.trainer.TrainerModule class method) (xlstm_jax.trainer.llm.trainer.LLMTrainer class method) load_huggingface_dataset() (in module xlstm_jax.dataset.grain_data_processing) load_model() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) load_model_params_and_config_from_checkpoint() (in module xlstm_jax.utils.model_param_handling.load) load_pretrained_model() (in module xlstm_jax.trainer.callbacks.checkpointing) (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) load_tokenizer() (in module xlstm_jax.dataset.hf_tokenizer) log_config() (xlstm_jax.trainer.logger.base_logger.Logger method) (xlstm_jax.trainer.logger.base_logger.LoggerTool method) (xlstm_jax.trainer.logger.file_logger.FileLogger method) (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLogger method) (xlstm_jax.trainer.logger.wandb_logger.WandBLogger method) log_dir (xlstm_jax.trainer.logger.base_logger.LoggerConfig property) (xlstm_jax.trainer.logger.file_logger.FileLoggerConfig attribute), [1] (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLoggerConfig attribute), [1] (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute), [1] log_epoch_key (xlstm_jax.trainer.logger.file_logger.FileLoggerConfig attribute), [1] log_every_n_steps (xlstm_jax.trainer.logger.base_logger.LoggerConfig attribute), [1] LOG_FORMAT (in module xlstm_jax.trainer.logger.cmd_logging) log_grad_norm (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) log_grad_norm_per_param (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) log_host_metrics() (xlstm_jax.trainer.logger.base_logger.Logger method) log_info() (in module lmeval_extended_evaluation) (in module xlstm_jax.train_init_fns) log_intermediates (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) log_logit_stats (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute), [1] log_lr_key (xlstm_jax.trainer.callbacks.lr_monitor.LearningRateMonitorConfig attribute) log_metrics() (xlstm_jax.trainer.logger.base_logger.LoggerTool method) (xlstm_jax.trainer.logger.file_logger.FileLogger method) (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLogger method) (xlstm_jax.trainer.logger.wandb_logger.WandBLogger method)

log_param_norm (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) log_param_norm_per_param (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) log_path (xlstm_jax.define_hydra_schemas.Config attribute) (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpointConfig attribute), [1] (xlstm_jax.trainer.callbacks.profiler.JaxProfiler attribute) (xlstm_jax.trainer.logger.base_logger.Logger attribute) (xlstm_jax.trainer.logger.base_logger.LoggerConfig attribute), [1] (xlstm_jax.trainer.logger.file_logger.FileLogger attribute) (xlstm_jax.trainer.logger.wandb_logger.WandBLogger attribute) log_step() (xlstm_jax.trainer.logger.base_logger.Logger method) log_step_key (xlstm_jax.trainer.logger.file_logger.FileLoggerConfig attribute), [1] log_tools (xlstm_jax.trainer.logger.base_logger.LoggerConfig attribute), [1] log_training_info() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) logcumsumexp_fwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) Logger (class in xlstm_jax.trainer.logger.base_logger) LOGGER (in module lmeval_extended_evaluation) (in module xlstm_jax.dataset.grain_batch_rampup) (in module xlstm_jax.dataset.grain_data_processing) (in module xlstm_jax.dataset.grain_iterator) (in module xlstm_jax.dataset.input_pipeline_interface) (in module xlstm_jax.dataset.lmeval_pipeline) (in module xlstm_jax.dataset.multihost_dataloading) (in module xlstm_jax.kernels) (in module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.layer_factory) (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.layer_factory) (in module xlstm_jax.models.xlstm_parallel.components.normalization) (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.layer_factory) (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.cuda_init) (in module xlstm_jax.train_init_fns) (in module xlstm_jax.trainer.base.trainer) (in module xlstm_jax.trainer.callbacks.checkpointing) (in module xlstm_jax.trainer.callbacks.extended_evaluation) (in module xlstm_jax.trainer.callbacks.lr_monitor) (in module xlstm_jax.trainer.callbacks.profiler) (in module xlstm_jax.trainer.logger.base_logger) (in module xlstm_jax.trainer.logger.cmd_logging) (in module xlstm_jax.trainer.logger.file_logger) (in module xlstm_jax.trainer.logger.tensorboard_logger) (in module xlstm_jax.trainer.logger.wandb_logger) (in module xlstm_jax.trainer.metrics) (in module xlstm_jax.trainer.optimizer.optimizer) (in module xlstm_jax.utils.model_param_handling.convert_state_dict) (in module xlstm_jax.utils.model_param_handling.handle_mlstm_simple) (in module xlstm_jax.utils.model_param_handling.load) (in module xlstm_jax.utils.model_param_handling.store) (in module xlstm_jax.utils.pytree_utils) logger (xlstm_jax.define_hydra_schemas.Config attribute) (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) (xlstm_jax.trainer.logger.file_logger.FileLogger attribute) (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLogger attribute) (xlstm_jax.trainer.logger.wandb_logger.WandBLogger attribute) LoggerConfig (class in xlstm_jax.trainer.logger.base_logger) loggers_to_use (xlstm_jax.define_hydra_schemas.BaseLoggerConfig attribute) LoggerTool (class in xlstm_jax.trainer.logger.base_logger) LoggerToolsConfig (class in xlstm_jax.trainer.logger.base_logger) logging_name (xlstm_jax.define_hydra_schemas.Config attribute) logits_soft_cap (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.shared.lm_head.TPLMHead attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) LogMode (in module xlstm_jax.common_types) logs (xlstm_jax.trainer.logger.file_logger.FileLogger attribute) loss_fn() (in module xlstm_jax.models.xlstm_parallel.training) loss_function() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) lr (xlstm_jax.define_hydra_schemas.Config attribute) (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig attribute), [1] lr_monitor (xlstm_jax.define_hydra_schemas.Config attribute) lr_scheduler (xlstm_jax.trainer.callbacks.lr_monitor.LearningRateMonitor attribute)

M

m1 (xlstm_jax.trainer.optimizer.ademamix.ScaleByAdemamixState attribute) m2 (xlstm_jax.trainer.optimizer.ademamix.ScaleByAdemamixState attribute) main_process_only (xlstm_jax.trainer.callbacks.callback.CallbackConfig attribute) (xlstm_jax.trainer.callbacks.lr_monitor.LearningRateMonitorConfig attribute) (xlstm_jax.trainer.callbacks.profiler.JaxProfilerConfig attribute), [1] main_train() (in module xlstm_jax.main_train) make_grain_iterator() (in module xlstm_jax.dataset.grain_data_processing) make_grain_llm_dataset() (in module xlstm_jax.dataset.grain_iterator) make_grain_llm_iterator() (in module xlstm_jax.dataset.grain_iterator) make_grain_multihost_iterator() (in module xlstm_jax.dataset.grain_iterator) manager (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint attribute) map() (xlstm_jax.dataset.grain_transforms.AddEODToken method) (xlstm_jax.dataset.grain_transforms.CollateToBatch method) (xlstm_jax.dataset.grain_transforms.HFNormalizeFeatures method) (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize method) (xlstm_jax.dataset.grain_transforms.HFTokenize method) (xlstm_jax.dataset.grain_transforms.InferSegmentations method) (xlstm_jax.dataset.grain_transforms.NormalizeFeatures method) (xlstm_jax.dataset.grain_transforms.PadToMaxLength method) (xlstm_jax.dataset.grain_transforms.ParseArrayRecords method) (xlstm_jax.dataset.grain_transforms.ParseFeatures method) (xlstm_jax.dataset.grain_transforms.ParseTokenizedArrayRecords method) (xlstm_jax.dataset.grain_transforms.ReformatLazyPacking static method) (xlstm_jax.dataset.grain_transforms.ReformatPacking static method) (xlstm_jax.dataset.grain_transforms.ShiftData method) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihood static method) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling method) (xlstm_jax.dataset.lmeval_pipeline.CompleteLLMIndexedBatch static method) (xlstm_jax.dataset.lmeval_pipeline.ParseLMEval method) mask_across_document_boundaries (xlstm_jax.models.llama.llama.LlamaConfig attribute) mask_except_model_idx (xlstm_jax.distributed.tensor_parallel.ModelParallelismWrapper attribute) max_length (xlstm_jax.dataset.grain_transforms.AddEODToken attribute) (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) (xlstm_jax.dataset.grain_transforms.HFTokenize attribute) (xlstm_jax.dataset.grain_transforms.PadToMaxLength attribute) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling attribute) max_length_prefix (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) max_steps_per_epoch (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) max_target_length (xlstm_jax.dataset.configs.DataConfig attribute) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) (xlstm_jax.dataset.configs.SyntheticDataConfig attribute) max_to_keep (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpointConfig attribute), [1] mesh (xlstm_jax.dataset.synthetic_dataloading.SyntheticDataIterator attribute) (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) metadata (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint attribute) metric_postprocess_fn (xlstm_jax.trainer.logger.base_logger.Logger attribute) MetricElement (in module xlstm_jax.common_types) Metrics (in module xlstm_jax.common_types) min_length (xlstm_jax.dataset.lmeval_pipeline.PadBatchDataset attribute) mlstm (xlstm_jax.models.xlstm_clean.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlockConfig attribute) mlstm_block (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) mlstm_cell (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) mlstm_chunkwise__parallel_bw_dK() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_bw_dK) mlstm_chunkwise__parallel_bw_dQ() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_bw_dQ) mlstm_chunkwise__parallel_bw_dV() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_bw_dV) mlstm_chunkwise__parallel_fw_Hintra() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_fw) mlstm_chunkwise__recurrent_bw_dC() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._recurrent_bw) mlstm_chunkwise__recurrent_fw_C() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._recurrent_fw) mlstm_chunkwise_bw() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._combined_bw) mlstm_chunkwise_fw() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._combined_fw) mlstm_chunkwise_max_triton() (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3.triton_fwbw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice.triton_fwbw) (in module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize.triton_fwbw) mlstm_chunkwise_triton_stablef() (in module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef.triton_fwbw) mlstm_fwbw_custom_grad() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw) mlstm_triton() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.triton_chunk) mLSTMBackend (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config) (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config) (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config) mLSTMBackendAttention (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention) mLSTMBackendAttentionConfig (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention) mLSTMBackendFwbw (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw) mLSTMBackendFwbwConfig (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw) mLSTMBackendJax (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.simple) mLSTMBackendJaxConfig (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.simple) mLSTMBackendNameAndKwargs (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.backend) (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend) (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend) mLSTMBackendParallel (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple) mLSTMBackendParallelConfig (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple) mLSTMBackendRecurrent (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent) mLSTMBackendRecurrentConfig (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent) mLSTMBackendRecurrentTriton (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton) mLSTMBackendRecurrentTritonConfig (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton) mLSTMBackendTorch (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple) mLSTMBackendTorchConfig (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple) mLSTMBackendTriton (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels) mLSTMBackendTritonConfig (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels) mLSTMBlock (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block) mLSTMBlock() (in module xlstm_jax.models.xlstm_clean.blocks.mlstm.block) (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.block) mLSTMBlockConfig (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.block) (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.block) (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block) mLSTMCell (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.cell) (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell) (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell) mLSTMCellConfig (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.cell) (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell) (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell) mLSTMFunc() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.triton_chunk) mLSTMFunction (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.triton_chunk) mLSTMfwbw (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw) mLSTMfwbwConfig (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw) mLSTMInnerLayer (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer) mLSTMLayer (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.layer) (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer) (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer) mLSTMLayerConfig (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.layer) (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer) (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer) mLSTMLayerV1 (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer_v1) mLSTMTorchFunction() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw) mod (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellCUDA attribute) mode (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpointConfig attribute), [1] (xlstm_jax.trainer.logger.base_logger.Logger property) mode_stack (xlstm_jax.trainer.logger.base_logger.Logger attribute) model (xlstm_jax.define_hydra_schemas.Config attribute) model_axis_name (xlstm_jax.distributed.pipeline_parallel.PipelineModule attribute) (xlstm_jax.distributed.tensor_parallel.ModelParallelismWrapper attribute) (xlstm_jax.distributed.tensor_parallel.TPAsyncDense attribute), [1] (xlstm_jax.distributed.tensor_parallel.TPDense attribute), [1] (xlstm_jax.models.configs.ParallelConfig attribute) model_axis_size (xlstm_jax.models.configs.ParallelConfig attribute) model_class (xlstm_jax.models.configs.ModelConfig attribute) model_config (xlstm_jax.models.configs.ModelConfig attribute) (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) ModelCheckpoint (class in xlstm_jax.trainer.callbacks.checkpointing) ModelCheckpointConfig (class in xlstm_jax.trainer.callbacks.checkpointing) ModelConfig (class in xlstm_jax.models.configs) ModelParallelismWrapper (class in xlstm_jax.distributed.tensor_parallel) module lmeval_extended_evaluation xlstm_jax xlstm_jax.common_types xlstm_jax.configs xlstm_jax.dataset xlstm_jax.dataset.batch xlstm_jax.dataset.configs xlstm_jax.dataset.grain_batch_rampup xlstm_jax.dataset.grain_data_processing xlstm_jax.dataset.grain_iterator xlstm_jax.dataset.grain_transforms xlstm_jax.dataset.hf_tokenizer xlstm_jax.dataset.input_pipeline_interface xlstm_jax.dataset.lmeval_dataset xlstm_jax.dataset.lmeval_pipeline xlstm_jax.dataset.multihost_dataloading xlstm_jax.dataset.synthetic_dataloading xlstm_jax.define_hydra_schemas xlstm_jax.distributed xlstm_jax.distributed.array_utils xlstm_jax.distributed.data_parallel xlstm_jax.distributed.mesh_utils xlstm_jax.distributed.pipeline_parallel xlstm_jax.distributed.single_gpu xlstm_jax.distributed.tensor_parallel xlstm_jax.distributed.xla_utils xlstm_jax.import_utils xlstm_jax.kernels xlstm_jax.kernels.kernel_utils xlstm_jax.kernels.mlstm_chunkwise xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3 xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_bw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_fw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3.triton_fwbw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_bw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_fw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice.triton_fwbw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._chunkwise_gates xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._combined_bw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._combined_fw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_bw_dK xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_bw_dQ xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_bw_dV xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_fw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._recurrent_bw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._recurrent_fw xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize.triton_fwbw xlstm_jax.kernels.mlstm_chunkwise.triton_stablef xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_bw xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_fw xlstm_jax.kernels.mlstm_chunkwise.triton_stablef.triton_fwbw xlstm_jax.kernels.mlstm_recurrent xlstm_jax.kernels.mlstm_recurrent.triton_fused_fw xlstm_jax.kernels.stride_utils xlstm_jax.main_train xlstm_jax.models xlstm_jax.models.configs xlstm_jax.models.llama xlstm_jax.models.llama.attention xlstm_jax.models.llama.feedforward xlstm_jax.models.llama.llama xlstm_jax.models.shared xlstm_jax.models.shared.init xlstm_jax.models.shared.lm_head xlstm_jax.models.shared.utils xlstm_jax.models.xlstm_clean xlstm_jax.models.xlstm_clean.blocks xlstm_jax.models.xlstm_clean.blocks.mlstm xlstm_jax.models.xlstm_clean.blocks.mlstm.backend xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config_utils xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.layer_factory xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.simple xlstm_jax.models.xlstm_clean.blocks.mlstm.block xlstm_jax.models.xlstm_clean.blocks.mlstm.cell xlstm_jax.models.xlstm_clean.blocks.mlstm.layer xlstm_jax.models.xlstm_clean.blocks.xlstm_block xlstm_jax.models.xlstm_clean.components xlstm_jax.models.xlstm_clean.components.conv xlstm_jax.models.xlstm_clean.components.feedforward xlstm_jax.models.xlstm_clean.components.init xlstm_jax.models.xlstm_clean.components.linear_headwise xlstm_jax.models.xlstm_clean.components.ln xlstm_jax.models.xlstm_clean.utils xlstm_jax.models.xlstm_clean.xlstm_block_stack xlstm_jax.models.xlstm_clean.xlstm_lm_model xlstm_jax.models.xlstm_parallel xlstm_jax.models.xlstm_parallel.benchmark xlstm_jax.models.xlstm_parallel.blocks xlstm_jax.models.xlstm_parallel.blocks.mlstm xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.layer_factory xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend_utils xlstm_jax.models.xlstm_parallel.blocks.mlstm.block xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer_v1 xlstm_jax.models.xlstm_parallel.blocks.xlstm_block xlstm_jax.models.xlstm_parallel.checkpointing xlstm_jax.models.xlstm_parallel.components xlstm_jax.models.xlstm_parallel.components.conv xlstm_jax.models.xlstm_parallel.components.feedforward xlstm_jax.models.xlstm_parallel.components.init xlstm_jax.models.xlstm_parallel.components.linear_headwise xlstm_jax.models.xlstm_parallel.components.normalization xlstm_jax.models.xlstm_parallel.training xlstm_jax.models.xlstm_parallel.utils xlstm_jax.models.xlstm_parallel.xlstm_block_stack xlstm_jax.models.xlstm_parallel.xlstm_lm_model xlstm_jax.models.xlstm_pytorch xlstm_jax.models.xlstm_pytorch.blocks xlstm_jax.models.xlstm_pytorch.blocks.mlstm xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config_utils xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.layer_factory xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.triton_chunk xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer xlstm_jax.models.xlstm_pytorch.blocks.slstm xlstm_jax.models.xlstm_pytorch.blocks.slstm.block xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer xlstm_jax.models.xlstm_pytorch.blocks.slstm.src xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.cuda_init xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla.lstm xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla.slstm xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block xlstm_jax.models.xlstm_pytorch.components xlstm_jax.models.xlstm_pytorch.components.conv xlstm_jax.models.xlstm_pytorch.components.feedforward xlstm_jax.models.xlstm_pytorch.components.init xlstm_jax.models.xlstm_pytorch.components.linear_headwise xlstm_jax.models.xlstm_pytorch.components.ln xlstm_jax.models.xlstm_pytorch.components.util xlstm_jax.models.xlstm_pytorch.utils xlstm_jax.models.xlstm_pytorch.xlstm_block_stack xlstm_jax.models.xlstm_pytorch.xlstm_lm_model xlstm_jax.resume_training xlstm_jax.start_training xlstm_jax.train_init_fns xlstm_jax.trainer xlstm_jax.trainer.base xlstm_jax.trainer.base.param_utils xlstm_jax.trainer.base.trainer xlstm_jax.trainer.callbacks xlstm_jax.trainer.callbacks.callback xlstm_jax.trainer.callbacks.checkpointing xlstm_jax.trainer.callbacks.extended_evaluation xlstm_jax.trainer.callbacks.lr_monitor xlstm_jax.trainer.callbacks.profiler xlstm_jax.trainer.data_module xlstm_jax.trainer.llm xlstm_jax.trainer.llm.sampling xlstm_jax.trainer.llm.trainer xlstm_jax.trainer.logger xlstm_jax.trainer.logger.base_logger xlstm_jax.trainer.logger.cmd_logging xlstm_jax.trainer.logger.file_logger xlstm_jax.trainer.logger.tensorboard_logger xlstm_jax.trainer.logger.wandb_logger xlstm_jax.trainer.metrics xlstm_jax.trainer.optimizer xlstm_jax.trainer.optimizer.ademamix xlstm_jax.trainer.optimizer.optimizer xlstm_jax.trainer.optimizer.scheduler xlstm_jax.utils xlstm_jax.utils.error_logging_utils xlstm_jax.utils.model_param_handling xlstm_jax.utils.model_param_handling.convert_checkpoint xlstm_jax.utils.model_param_handling.convert_state_dict xlstm_jax.utils.model_param_handling.handle_mlstm_simple xlstm_jax.utils.model_param_handling.load xlstm_jax.utils.model_param_handling.store xlstm_jax.utils.pytree_utils

module (xlstm_jax.models.xlstm_pytorch.components.util.ParameterProxy attribute) module_fn (xlstm_jax.distributed.pipeline_parallel.PipelineModule attribute) (xlstm_jax.distributed.tensor_parallel.ModelParallelismWrapper attribute) module_fns (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.ResidualBlock attribute) module_kwargs (xlstm_jax.distributed.tensor_parallel.ModelParallelismWrapper attribute) monitor (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpointConfig attribute), [1] move_mlstm_jax_state_dict_into_torch_state_dict() (in module xlstm_jax.utils.model_param_handling.handle_mlstm_simple) move_safetensors_state_dict_params_() (in module xlstm_jax.utils.model_param_handling.convert_state_dict) move_state_dict_params_() (in module xlstm_jax.utils.model_param_handling.convert_state_dict) multihead_attention() (in module xlstm_jax.models.llama.attention) MultiHeadLayerNorm (class in xlstm_jax.models.xlstm_pytorch.components.ln) MultiHeadLayerNorm() (in module xlstm_jax.models.xlstm_clean.components.ln) MultiHeadNormLayer() (in module xlstm_jax.models.xlstm_parallel.components.normalization) MultiHostDataLoadIterator (class in xlstm_jax.dataset.multihost_dataloading) MultihostSortedRemapDataset (class in xlstm_jax.dataset.lmeval_pipeline) multiple_of (xlstm_jax.dataset.lmeval_pipeline.PadBatchDataset attribute) (xlstm_jax.dataset.lmeval_pipeline.PadSequenceInBatchDataset attribute) (xlstm_jax.models.llama.feedforward.FeedForwardConfig attribute) mutable_variables (xlstm_jax.common_types.TrainState attribute) (xlstm_jax.trainer.callbacks.extended_evaluation.EvalState attribute) MutableMetricElement (in module xlstm_jax.common_types) MutableMetrics (in module xlstm_jax.common_types)

N

n_gpus (xlstm_jax.define_hydra_schemas.Config attribute) name (xlstm_jax.dataset.configs.DataConfig attribute) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) (xlstm_jax.dataset.configs.SyntheticDataConfig attribute) (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config_utils.NameAndKwargs attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.mLSTMBackendNameAndKwargs attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils.NameAndKwargs attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.mLSTMBackendNameAndKwargs attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config_utils.NameAndKwargs attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.mLSTMBackendNameAndKwargs attribute) (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig attribute), [1] NameAndKwargs (class in xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config_utils) (class in xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils) (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config_utils) ndim (xlstm_jax.models.xlstm_pytorch.components.ln.LayerNorm attribute) (xlstm_jax.models.xlstm_pytorch.components.ln.MultiHeadLayerNorm attribute) (xlstm_jax.models.xlstm_pytorch.components.util.ParameterProxy property) nesterov (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] norm_eps (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) norm_fn (xlstm_jax.models.shared.lm_head.TPLMHead attribute) norm_type (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) norm_type_v1 (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) norm_val (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTritonConfig attribute)

NormalizeFeatures (class in xlstm_jax.dataset.grain_transforms) NormLayer() (in module xlstm_jax.models.xlstm_parallel.components.normalization) NormType (in module xlstm_jax.models.xlstm_parallel.components.normalization) nu (xlstm_jax.trainer.optimizer.ademamix.ScaleByAdemamixState attribute) num_batches (xlstm_jax.dataset.configs.SyntheticDataConfig attribute) (xlstm_jax.dataset.synthetic_dataloading.SyntheticDataIterator attribute) num_blocks (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) num_epochs (xlstm_jax.define_hydra_schemas.Config attribute) num_fewshot (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) num_gates (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) num_heads (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpandConfig attribute) num_layers (xlstm_jax.models.llama.attention.SelfAttentionConfig attribute) (xlstm_jax.models.llama.feedforward.FeedForwardConfig attribute) num_microbatches (xlstm_jax.distributed.pipeline_parallel.PipelineModule attribute) num_states (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) num_train_steps (xlstm_jax.define_hydra_schemas.Config attribute)

O

ogate (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer attribute) ogate_act_fn (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) on_extended_evaluation_end() (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) on_extended_evaluation_start() (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation method) on_filtered_training_epoch_end() (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.callbacks.lr_monitor.LearningRateMonitor method) on_filtered_training_epoch_start() (xlstm_jax.trainer.callbacks.callback.Callback method) on_filtered_training_step() (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.callbacks.lr_monitor.LearningRateMonitor method) on_filtered_validation_epoch_end() (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint method) on_filtered_validation_epoch_start() (lmeval_extended_evaluation.LMEvalEvaluation method) (xlstm_jax.trainer.callbacks.callback.Callback method) on_test_epoch_end() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) on_test_epoch_start() (lmeval_extended_evaluation.LMEvalEvaluation method) (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) on_training_end() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) on_training_epoch_end() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.callbacks.profiler.JaxProfiler method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method)

on_training_epoch_start() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) on_training_start() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.callbacks.profiler.JaxProfiler method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) (xlstm_jax.trainer.logger.base_logger.Logger method) on_training_step() (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.callbacks.profiler.JaxProfiler method) on_validation_epoch_end() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) on_validation_epoch_start() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.callbacks.profiler.JaxProfiler method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) optimizer (xlstm_jax.define_hydra_schemas.Config attribute) optimizer_config (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) OptimizerConfig (class in xlstm_jax.trainer.optimizer.optimizer) outnorm (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell attribute) output_init_fn (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) output_shape (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute)

P

P (in module xlstm_jax.utils.error_logging_utils) pad_dataset() (in module xlstm_jax.dataset.grain_data_processing) pad_elem (xlstm_jax.dataset.lmeval_pipeline.PadBatchDataset attribute) pad_shapes (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator attribute) pad_token_id (xlstm_jax.dataset.grain_transforms.ShiftData attribute) pad_value (xlstm_jax.dataset.lmeval_pipeline.PadSequenceInBatchDataset attribute) (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator attribute) PadBatchDataset (class in xlstm_jax.dataset.lmeval_pipeline) PaddedDataset (class in xlstm_jax.dataset.grain_data_processing) PadSequenceInBatchDataset (class in xlstm_jax.dataset.lmeval_pipeline) PadToMaxLength (class in xlstm_jax.dataset.grain_transforms) parallel (xlstm_jax.define_hydra_schemas.Config attribute) (xlstm_jax.models.configs.ModelConfig attribute) (xlstm_jax.models.llama.attention.SelfAttentionConfig attribute) (xlstm_jax.models.llama.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.shared.lm_head.TPLMHead attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.block.mLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) parallel_stabilized_simple() (in module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.simple) (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple) (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple) ParallelConfig (class in xlstm_jax.models.configs) Parameter (in module xlstm_jax.common_types) parameter_name (xlstm_jax.models.xlstm_pytorch.components.util.ParameterProxy attribute) ParameterProxy (class in xlstm_jax.models.xlstm_pytorch.components.util) parameters_to_dtype() (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) params (xlstm_jax.trainer.callbacks.extended_evaluation.EvalState attribute) ParseArrayRecords (class in xlstm_jax.dataset.grain_transforms) ParseFeatures (class in xlstm_jax.dataset.grain_transforms) ParseLMEval (class in xlstm_jax.dataset.lmeval_pipeline) ParseTokenizedArrayRecords (class in xlstm_jax.dataset.grain_transforms)

pipeline_axis_name (xlstm_jax.models.configs.ParallelConfig attribute) pipeline_axis_size (xlstm_jax.models.configs.ParallelConfig attribute) pipeline_convert_mlstm_checkpoint_jax_to_torch_simple() (in module xlstm_jax.utils.model_param_handling.handle_mlstm_simple) PipelineModule (class in xlstm_jax.distributed.pipeline_parallel) pointwise (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla attribute) precompute_freqs() (in module xlstm_jax.models.llama.attention) (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention) prefix_column_name (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) prefix_tokenizer (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) prepare_module() (in module xlstm_jax.models.shared.utils) preprocess_dataset() (in module xlstm_jax.dataset.grain_data_processing) print_metrics() (in module xlstm_jax.models.xlstm_parallel.training) PRNGKeyArray (in module xlstm_jax.common_types) profile_every_n_minutes (xlstm_jax.trainer.callbacks.profiler.JaxProfiler attribute) (xlstm_jax.trainer.callbacks.profiler.JaxProfilerConfig attribute), [1] profile_first_step (xlstm_jax.trainer.callbacks.profiler.JaxProfiler attribute) (xlstm_jax.trainer.callbacks.profiler.JaxProfilerConfig attribute), [1] profile_log_dir (xlstm_jax.trainer.callbacks.profiler.JaxProfilerConfig attribute), [1] profile_n_steps (xlstm_jax.trainer.callbacks.profiler.JaxProfiler attribute) (xlstm_jax.trainer.callbacks.profiler.JaxProfilerConfig attribute), [1] profiler_active (xlstm_jax.trainer.callbacks.profiler.JaxProfiler attribute) profiler_last_time (xlstm_jax.trainer.callbacks.profiler.JaxProfiler attribute) profiling (xlstm_jax.define_hydra_schemas.Config attribute) proj_down (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.GatedFeedForward attribute) proj_factor (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_clean.utils.UpProjConfigMixin attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.utils.UpProjConfigMixin attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.utils.UpProjConfigMixin attribute) proj_up (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.GatedFeedForward attribute) PyTree (in module xlstm_jax.common_types) pytree_diff() (in module xlstm_jax.utils.pytree_utils) pytree_key_path_to_str() (in module xlstm_jax.utils.pytree_utils)

Q

q_proj (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) qk_dim_factor (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) qk_norm (xlstm_jax.models.llama.attention.SelfAttentionConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute)

qk_pre_activation_function (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention.mLSTMBackendAttentionConfig attribute) qkv_proj_blocksize (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute)

R

R (in module xlstm_jax.utils.error_logging_utils) raw_generate_synthetic_data() (xlstm_jax.dataset.synthetic_dataloading.SyntheticDataIterator static method) recurrent_cumsum_bwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) recurrent_cumsum_fwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) recurrent_kernel (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase attribute) recurrent_reversed_cumsum_bwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) recurrent_reversed_cumsum_fwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) recurrent_sequence_fw() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent) recurrent_step_fw() (in module xlstm_jax.kernels.mlstm_recurrent.triton_fused_fw) (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent) recurrent_step_stabilized_simple() (in module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.simple) (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple) recurrent_weight_init (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig attribute) RecursionLimit (class in xlstm_jax.utils.pytree_utils) reduce_slicing (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTritonConfig attribute) ReformatLazyPacking (class in xlstm_jax.dataset.grain_transforms) ReformatPacking (class in xlstm_jax.dataset.grain_transforms) register_configs() (in module xlstm_jax.define_hydra_schemas) remat (xlstm_jax.models.configs.ParallelConfig attribute) request_name (xlstm_jax.dataset.lmeval_pipeline.ParseLMEval attribute) reset() (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator method) reset_after_epoch (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator attribute) reset_at_document_boundaries (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) reset_fgate_value (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig attribute) reset_parameters() (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple.mLSTMBackendTorch method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlock method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlock method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer method) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlock method) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1d method) (xlstm_jax.models.xlstm_pytorch.components.feedforward.GatedFeedForward method) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpand method) (xlstm_jax.models.xlstm_pytorch.components.ln.LayerNorm method) (xlstm_jax.models.xlstm_pytorch.components.ln.MultiHeadLayerNorm method) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStack method) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel method)

residual_weight (xlstm_jax.models.xlstm_pytorch.components.ln.LayerNorm attribute) (xlstm_jax.models.xlstm_pytorch.components.ln.MultiHeadLayerNorm attribute) ResidualBlock (class in xlstm_jax.models.xlstm_parallel.blocks.xlstm_block) resolve_import() (in module xlstm_jax.import_utils) resolve_import_from_string() (in module xlstm_jax.import_utils) resolve_norm() (in module xlstm_jax.models.xlstm_parallel.components.normalization) resolve_step_idx() (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint method) restore_data_loaders() (xlstm_jax.trainer.base.trainer.TrainerModule static method) (xlstm_jax.trainer.llm.trainer.LLMTrainer static method) restore_model() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) resume_training() (in module xlstm_jax.resume_training) return_state (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbwConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) rev_cumsum() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw) (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw) rev_cumsum_off() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw) (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw) reversed_cumsum() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) ReversedCumsumFunction (class in xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) rng (xlstm_jax.common_types.TrainState attribute) (xlstm_jax.trainer.callbacks.extended_evaluation.EvalState attribute) rnn_function_registry (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) round_proj_up_dim_up (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_clean.utils.UpProjConfigMixin attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.utils.UpProjConfigMixin attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.utils.UpProjConfigMixin attribute) round_proj_up_to_multiple_of (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_clean.utils.UpProjConfigMixin attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_parallel.utils.UpProjConfigMixin attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.xlstm_pytorch.utils.UpProjConfigMixin attribute) round_to_multiple() (in module xlstm_jax.models.xlstm_pytorch.components.util) run_backend() (in module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend_utils) run_evaluate() (lmeval_extended_evaluation.LMEvalEvaluation method) run_model_init() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method)

S

S (in module xlstm_jax.dataset.grain_batch_rampup) save_checkpoint() (in module xlstm_jax.models.xlstm_parallel.checkpointing) save_dataloader() (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint method) save_dataloader_state (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpointConfig attribute), [1] save_model() (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpoint method) save_optimizer_state (xlstm_jax.trainer.callbacks.checkpointing.ModelCheckpointConfig attribute), [1] scale (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) scale_by_ademamix() (in module xlstm_jax.trainer.optimizer.ademamix) scale_init() (in module xlstm_jax.distributed.tensor_parallel) ScaleByAdemamixState (class in xlstm_jax.trainer.optimizer.ademamix) scan_blocks (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) scheduler (xlstm_jax.define_hydra_schemas.Config attribute) (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] SchedulerConfig (class in xlstm_jax.trainer.optimizer.scheduler) seed (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) seed_eval (xlstm_jax.trainer.base.trainer.TrainerConfig attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig attribute) segment_mask() (in module xlstm_jax.models.llama.attention) SelfAttention (class in xlstm_jax.models.llama.attention) SelfAttentionBlock (class in xlstm_jax.models.llama.llama) SelfAttentionConfig (class in xlstm_jax.models.llama.attention) sequence_idx (xlstm_jax.dataset.batch.LLMIndexedBatch attribute) sequence_to_bytestring() (xlstm_jax.dataset.grain_transforms.ParseTokenizedArrayRecords static method) set_dataset() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.callbacks.callback.Callback method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) set_state() (xlstm_jax.dataset.grain_batch_rampup._BatchRampUpDatasetIterator method) (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator method) set_XLA_flags() (in module xlstm_jax.distributed.xla_utils) setup() (xlstm_jax.trainer.logger.file_logger.FileLogger method) (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLogger method) (xlstm_jax.trainer.logger.wandb_logger.WandBLogger method) setup_exception_logging() (in module xlstm_jax.trainer.logger.cmd_logging) setup_logging() (in module xlstm_jax.trainer.logger.cmd_logging) setup_logging_multiprocess() (in module xlstm_jax.trainer.logger.cmd_logging) shape (xlstm_jax.models.xlstm_pytorch.components.util.ParameterProxy property) shard_grain_dataset() (in module xlstm_jax.dataset.grain_iterator) shard_module_params() (in module xlstm_jax.distributed.data_parallel) shard_params() (in module xlstm_jax.distributed.data_parallel) shift_and_refine() (in module xlstm_jax.dataset.grain_transforms) shift_left() (in module xlstm_jax.dataset.grain_transforms) shift_right() (in module xlstm_jax.dataset.grain_transforms) shift_target (xlstm_jax.dataset.grain_transforms.ShiftData attribute) ShiftData (class in xlstm_jax.dataset.grain_transforms) shuffle_data (xlstm_jax.dataset.configs.DataConfig attribute) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) (xlstm_jax.dataset.configs.SyntheticDataConfig attribute) simple_array() (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling method) simulate_CPU_devices() (in module xlstm_jax.distributed.xla_utils) skip_communication (xlstm_jax.distributed.tensor_parallel.TPDense attribute), [1] slstm (xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.xlstm_block.xLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlockConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlockConfig attribute) slstm_at (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) slstm_block (xlstm_jax.models.xlstm_clean.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStackConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute)

slstm_cell (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer attribute) slstm_forward() (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla) slstm_forward_pointwise() (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla.lstm) (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla.slstm) slstm_forward_step() (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla) slstm_pointwise_function_registry (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla) sLSTMBlock (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.block) sLSTMBlockConfig (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.block) sLSTMCell (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) sLSTMCell_cuda (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) sLSTMCell_vanilla (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) sLSTMCellBase (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) sLSTMCellConfig (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) sLSTMCellCUDA (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) sLSTMCellFuncGenerator() (in module xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell) sLSTMLayer (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer) sLSTMLayerConfig (class in xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer) small_init() (in module xlstm_jax.models.shared.init) (in module xlstm_jax.models.xlstm_clean.components.init) small_init_init_() (in module xlstm_jax.models.xlstm_pytorch.components.init) soft_cap_logits() (in module xlstm_jax.models.shared.utils) softmax_bwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) softmax_fwd_kernel() (in module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils) SortedDataset (class in xlstm_jax.dataset.lmeval_pipeline) split (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) split_array_over_mesh() (in module xlstm_jax.distributed.array_utils) split_rngs (xlstm_jax.distributed.tensor_parallel.ModelParallelismWrapper attribute) stabilize_correctly (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbwConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels.mLSTMBackendTritonConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) stack_params() (in module xlstm_jax.distributed.array_utils) start_epoch() (xlstm_jax.trainer.logger.base_logger.Logger method) start_trace() (xlstm_jax.trainer.callbacks.profiler.JaxProfiler method) start_training() (in module xlstm_jax.start_training) state (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) state_dtype (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent.mLSTMBackendRecurrentConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton.mLSTMBackendRecurrentTritonConfig attribute) state_set (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator attribute) step (xlstm_jax.trainer.callbacks.extended_evaluation.EvalState attribute) (xlstm_jax.trainer.logger.base_logger.Logger attribute) step() (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlock method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell.mLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlock method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_cuda method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCell_vanilla method) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellBase method) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlock method) (xlstm_jax.models.xlstm_pytorch.components.conv.CausalConv1d method) (xlstm_jax.models.xlstm_pytorch.xlstm_block_stack.xLSTMBlockStack method) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel method) step_counter (xlstm_jax.dataset.multihost_dataloading.MultiHostDataLoadIterator attribute) StepMetrics (in module xlstm_jax.common_types) StepMetricsElement (in module xlstm_jax.common_types) stepwise_rampup_schedule() (in module xlstm_jax.dataset.grain_batch_rampup) stop_trace() (xlstm_jax.trainer.callbacks.profiler.JaxProfiler method) store_checkpoint_sharded() (in module xlstm_jax.utils.model_param_handling.store) store_mlstm_simple_to_checkpoint() (in module xlstm_jax.utils.model_param_handling.handle_mlstm_simple) SubModelConfig (class in xlstm_jax.models.configs) sync_gradients() (in module xlstm_jax.distributed.data_parallel) SyntheticDataConfig (class in xlstm_jax.dataset.configs) SyntheticDataIterator (class in xlstm_jax.dataset.synthetic_dataloading) system_instruction (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute)

T

T (in module xlstm_jax.dataset.configs) (in module xlstm_jax.dataset.grain_batch_rampup) tabulate_params() (in module xlstm_jax.models.xlstm_parallel.training) (in module xlstm_jax.trainer.base.param_utils) (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) targets (xlstm_jax.dataset.batch.Batch attribute) (xlstm_jax.dataset.batch.LLMBatch attribute) (xlstm_jax.dataset.batch.LLMIndexedBatch attribute) targets_position (xlstm_jax.dataset.batch.LLMBatch attribute) (xlstm_jax.dataset.batch.LLMIndexedBatch attribute) targets_segmentation (xlstm_jax.dataset.batch.LLMBatch attribute) (xlstm_jax.dataset.batch.LLMIndexedBatch attribute) task_name (xlstm_jax.define_hydra_schemas.Config attribute) tb_flush_secs (xlstm_jax.define_hydra_schemas.BaseLoggerConfig attribute) (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLoggerConfig attribute), [1] tb_log_dir (xlstm_jax.define_hydra_schemas.BaseLoggerConfig attribute) tb_max_queue (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLoggerConfig attribute), [1] tb_new_style (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLoggerConfig attribute), [1] temperature_sampling() (in module xlstm_jax.trainer.llm.sampling) TensorBoardLogger (class in xlstm_jax.trainer.logger.tensorboard_logger) TensorBoardLoggerConfig (class in xlstm_jax.trainer.logger.tensorboard_logger) test_dataloader (xlstm_jax.trainer.data_module.DataloaderModule attribute) test_eval_function() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) test_model() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) text_column_name (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) theta (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention.mLSTMBackendAttentionConfig attribute) tie_weights (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) to_dict() (xlstm_jax.configs.ConfigDict method) (xlstm_jax.models.configs.ModelConfig method) (xlstm_jax.models.configs.SubModelConfig method) (xlstm_jax.models.llama.attention.SelfAttentionConfig method) (xlstm_jax.models.llama.feedforward.FeedForwardConfig method) (xlstm_jax.models.llama.llama.LlamaConfig method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils.NameAndKwargs method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.mLSTMBackendNameAndKwargs method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.block.mLSTMBlockConfig method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell.mLSTMCellConfig method) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig method) (xlstm_jax.models.xlstm_parallel.components.feedforward.FeedForwardConfig method) (xlstm_jax.models.xlstm_parallel.utils.UpProjConfigMixin method) (xlstm_jax.models.xlstm_parallel.xlstm_block_stack.xLSTMBlockStackConfig method) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig method) (xlstm_jax.trainer.base.trainer.TrainerConfig method) (xlstm_jax.trainer.callbacks.callback.CallbackConfig method) (xlstm_jax.trainer.llm.trainer.LLMTrainerConfig method) (xlstm_jax.trainer.logger.base_logger.LoggerConfig method) (xlstm_jax.trainer.logger.base_logger.LoggerToolsConfig method) (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig method) (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig method)

token_embedding (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel attribute) token_length() (in module xlstm_jax.dataset.lmeval_pipeline) tokenize (xlstm_jax.dataset.grain_transforms.NormalizeFeatures attribute) tokenize_data (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) tokenizer (xlstm_jax.dataset.grain_transforms.HFPrefixTokenize attribute) (xlstm_jax.dataset.grain_transforms.HFTokenize attribute) (xlstm_jax.dataset.lmeval_dataset.HFTokenizeLogLikelihoodRolling attribute) tokenizer_path (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) torch_dtype (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig property) torch_dtype_b (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig property) torch_dtype_r (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig property) torch_dtype_s (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig property) torch_dtype_w (xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell.sLSTMCellConfig property) tp_async_dense (xlstm_jax.models.configs.ParallelConfig attribute) tp_mode (xlstm_jax.distributed.tensor_parallel.TPAsyncDense attribute), [1] (xlstm_jax.distributed.tensor_parallel.TPDense attribute), [1] TPAsyncDense (class in xlstm_jax.distributed.tensor_parallel) TPDense (class in xlstm_jax.distributed.tensor_parallel) TPLMHead (class in xlstm_jax.models.shared.lm_head) tracker() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) train (xlstm_jax.models.llama.llama.FFNBlock attribute) (xlstm_jax.models.llama.llama.SelfAttentionBlock attribute) (xlstm_jax.models.llama.llama.TransformerBlock attribute) train_dataloader (xlstm_jax.trainer.data_module.DataloaderModule attribute) train_model() (xlstm_jax.trainer.base.trainer.TrainerModule method) (xlstm_jax.trainer.llm.trainer.LLMTrainer method) train_step() (in module xlstm_jax.models.xlstm_parallel.training) trainable_bias (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpandConfig attribute) trainable_weight (xlstm_jax.models.xlstm_clean.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_parallel.components.linear_headwise.LinearHeadwiseExpandConfig attribute) (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpandConfig attribute) trainer (xlstm_jax.define_hydra_schemas.Config attribute) (xlstm_jax.trainer.callbacks.callback.Callback attribute) (xlstm_jax.trainer.callbacks.extended_evaluation.ExtendedEvaluation attribute) trainer_config (xlstm_jax.trainer.base.trainer.TrainerModule attribute) (xlstm_jax.trainer.llm.trainer.LLMTrainer attribute) TrainerConfig (class in xlstm_jax.trainer.base.trainer) TrainerModule (class in xlstm_jax.trainer.base.trainer) TrainState (class in xlstm_jax.common_types) TransformerBlock (class in xlstm_jax.models.llama.llama) TransformerBlockStack (class in xlstm_jax.models.llama.llama) tree_bias_correction() (in module xlstm_jax.trainer.optimizer.ademamix) tree_cast() (in module xlstm_jax.trainer.optimizer.ademamix) tree_update_moment() (in module xlstm_jax.trainer.optimizer.ademamix) tree_update_moment_per_elem_norm() (in module xlstm_jax.trainer.optimizer.ademamix) tree_zeros_like() (in module xlstm_jax.trainer.optimizer.ademamix)

U

uniform_init() (in module xlstm_jax.models.shared.init) (in module xlstm_jax.models.xlstm_clean.components.init) unstack_params() (in module xlstm_jax.distributed.array_utils) update_metrics() (in module xlstm_jax.trainer.metrics) UpProjConfigMixin (class in xlstm_jax.models.xlstm_clean.utils) (class in xlstm_jax.models.xlstm_parallel.utils) (class in xlstm_jax.models.xlstm_pytorch.utils) use_bias (xlstm_jax.distributed.tensor_parallel.TPAsyncDense attribute), [1] (xlstm_jax.distributed.tensor_parallel.TPDense attribute), [1] (xlstm_jax.models.llama.attention.SelfAttentionConfig attribute) (xlstm_jax.models.llama.feedforward.FeedForwardConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute)

use_bidirectional_gather (xlstm_jax.distributed.tensor_parallel.TPAsyncDense attribute), [1] use_bidirectional_scatter (xlstm_jax.distributed.tensor_parallel.TPAsyncDense attribute), [1] use_infinite_eval (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) use_initial_state (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw.mLSTMBackendFwbwConfig attribute) (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw.mLSTMfwbwConfig attribute) use_scan (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent.mLSTMBackendRecurrentConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton.mLSTMBackendRecurrentTritonConfig attribute) use_sharded_clip_norm (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1]

V

v_dim_factor (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute) v_proj (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayer attribute) val_dataloader (xlstm_jax.trainer.data_module.DataloaderModule attribute) vmap_qk (xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig attribute) (xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer.mLSTMLayerConfig attribute)

vocab_size (xlstm_jax.define_hydra_schemas.CombinedModelConfig attribute) (xlstm_jax.models.llama.llama.LlamaConfig attribute) (xlstm_jax.models.shared.lm_head.TPLMHead attribute) (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute)

W

wandb_run (xlstm_jax.trainer.logger.wandb_logger.WandBLogger attribute) WandBLogger (class in xlstm_jax.trainer.logger.wandb_logger) WandBLoggerConfig (class in xlstm_jax.trainer.logger.wandb_logger) wang_init() (in module xlstm_jax.models.shared.init) (in module xlstm_jax.models.xlstm_clean.components.init) wang_init_() (in module xlstm_jax.models.xlstm_pytorch.components.init) warmup_steps (xlstm_jax.trainer.optimizer.scheduler.SchedulerConfig attribute), [1] wb_entity (xlstm_jax.define_hydra_schemas.BaseLoggerConfig attribute) (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute), [1] wb_host (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute), [1] wb_key (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute), [1] wb_name (xlstm_jax.define_hydra_schemas.BaseLoggerConfig attribute) (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute), [1] wb_notes (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute), [1] wb_project (xlstm_jax.define_hydra_schemas.BaseLoggerConfig attribute) (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute), [1] wb_resume_id (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute) wb_settings (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute), [1] wb_tags (xlstm_jax.define_hydra_schemas.BaseLoggerConfig attribute) (xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig attribute), [1] weight (xlstm_jax.models.xlstm_pytorch.components.linear_headwise.LinearHeadwiseExpand attribute) (xlstm_jax.models.xlstm_pytorch.components.ln.LayerNorm attribute) (xlstm_jax.models.xlstm_pytorch.components.ln.MultiHeadLayerNorm attribute) weight1 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) weight10 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute)

weight2 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) weight3 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) weight4 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) weight5 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) weight6 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) weight7 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) weight8 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) weight9 (xlstm_jax.define_hydra_schemas.DataTrainConfig attribute) weight_decay (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] weight_decay_exclude (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] weight_decay_include (xlstm_jax.trainer.optimizer.optimizer.OptimizerConfig attribute), [1] weight_decay_on_embedding (xlstm_jax.models.xlstm_clean.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_parallel.xlstm_lm_model.xLSTMLMModelConfig attribute) (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModelConfig attribute) weight_proxy (xlstm_jax.models.xlstm_pytorch.components.ln.LayerNorm property) (xlstm_jax.models.xlstm_pytorch.components.ln.MultiHeadLayerNorm property) WeightDecayOptimGroupMixin (class in xlstm_jax.models.xlstm_pytorch.utils) with_error_handling() (in module xlstm_jax.utils.error_logging_utils) worker_buffer_size (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) worker_count (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) (xlstm_jax.dataset.configs.GrainArrayRecordsDataConfig attribute) (xlstm_jax.dataset.configs.HFHubDataConfig attribute) write_out (lmeval_extended_evaluation.LMEvalEvaluationConfig attribute) writer (xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLogger attribute)

X

xlstm_block_stack (xlstm_jax.models.xlstm_pytorch.xlstm_lm_model.xLSTMLMModel attribute) xlstm_jax module xlstm_jax.common_types module xlstm_jax.configs module xlstm_jax.dataset module xlstm_jax.dataset.batch module xlstm_jax.dataset.configs module xlstm_jax.dataset.grain_batch_rampup module xlstm_jax.dataset.grain_data_processing module xlstm_jax.dataset.grain_iterator module xlstm_jax.dataset.grain_transforms module xlstm_jax.dataset.hf_tokenizer module xlstm_jax.dataset.input_pipeline_interface module xlstm_jax.dataset.lmeval_dataset module xlstm_jax.dataset.lmeval_pipeline module xlstm_jax.dataset.multihost_dataloading module xlstm_jax.dataset.synthetic_dataloading module xlstm_jax.define_hydra_schemas module xlstm_jax.distributed module xlstm_jax.distributed.array_utils module xlstm_jax.distributed.data_parallel module xlstm_jax.distributed.mesh_utils module xlstm_jax.distributed.pipeline_parallel module xlstm_jax.distributed.single_gpu module xlstm_jax.distributed.tensor_parallel module xlstm_jax.distributed.xla_utils module xlstm_jax.import_utils module xlstm_jax.kernels module xlstm_jax.kernels.kernel_utils module xlstm_jax.kernels.mlstm_chunkwise module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3 module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_bw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3._triton_fw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3.triton_fwbw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_bw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice._triton_fw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v3noslice.triton_fwbw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._chunkwise_gates module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._combined_bw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._combined_fw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_bw_dK module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_bw_dQ module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_bw_dV module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._parallel_fw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._recurrent_bw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize._recurrent_fw module xlstm_jax.kernels.mlstm_chunkwise.max_triton_fwbw_v5xlchunksize.triton_fwbw module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_bw module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef._triton_fw module xlstm_jax.kernels.mlstm_chunkwise.triton_stablef.triton_fwbw module xlstm_jax.kernels.mlstm_recurrent module xlstm_jax.kernels.mlstm_recurrent.triton_fused_fw module xlstm_jax.kernels.stride_utils module xlstm_jax.main_train module xlstm_jax.models module xlstm_jax.models.configs module xlstm_jax.models.llama module xlstm_jax.models.llama.attention module xlstm_jax.models.llama.feedforward module xlstm_jax.models.llama.llama module xlstm_jax.models.shared module xlstm_jax.models.shared.init module xlstm_jax.models.shared.lm_head module xlstm_jax.models.shared.utils module xlstm_jax.models.xlstm_clean module xlstm_jax.models.xlstm_clean.blocks module xlstm_jax.models.xlstm_clean.blocks.mlstm module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.config_utils module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.layer_factory module xlstm_jax.models.xlstm_clean.blocks.mlstm.backend.simple module xlstm_jax.models.xlstm_clean.blocks.mlstm.block module xlstm_jax.models.xlstm_clean.blocks.mlstm.cell module xlstm_jax.models.xlstm_clean.blocks.mlstm.layer module xlstm_jax.models.xlstm_clean.blocks.xlstm_block module xlstm_jax.models.xlstm_clean.components module xlstm_jax.models.xlstm_clean.components.conv module xlstm_jax.models.xlstm_clean.components.feedforward module xlstm_jax.models.xlstm_clean.components.init module xlstm_jax.models.xlstm_clean.components.linear_headwise module xlstm_jax.models.xlstm_clean.components.ln module xlstm_jax.models.xlstm_clean.utils module xlstm_jax.models.xlstm_clean.xlstm_block_stack module xlstm_jax.models.xlstm_clean.xlstm_lm_model module xlstm_jax.models.xlstm_parallel module xlstm_jax.models.xlstm_parallel.benchmark module xlstm_jax.models.xlstm_parallel.blocks module xlstm_jax.models.xlstm_parallel.blocks.mlstm module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.layer_factory module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple module

xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels module xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend_utils module xlstm_jax.models.xlstm_parallel.blocks.mlstm.block module xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell module xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer module xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer_v1 module xlstm_jax.models.xlstm_parallel.blocks.xlstm_block module xlstm_jax.models.xlstm_parallel.checkpointing module xlstm_jax.models.xlstm_parallel.components module xlstm_jax.models.xlstm_parallel.components.conv module xlstm_jax.models.xlstm_parallel.components.feedforward module xlstm_jax.models.xlstm_parallel.components.init module xlstm_jax.models.xlstm_parallel.components.linear_headwise module xlstm_jax.models.xlstm_parallel.components.normalization module xlstm_jax.models.xlstm_parallel.training module xlstm_jax.models.xlstm_parallel.utils module xlstm_jax.models.xlstm_parallel.xlstm_block_stack module xlstm_jax.models.xlstm_parallel.xlstm_lm_model module xlstm_jax.models.xlstm_pytorch module xlstm_jax.models.xlstm_pytorch.blocks module xlstm_jax.models.xlstm_pytorch.blocks.mlstm module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.config_utils module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.fwbw module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.layer_factory module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.simple module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.tl_utils module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.backend.triton_chunk module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.cell module xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer module xlstm_jax.models.xlstm_pytorch.blocks.slstm module xlstm_jax.models.xlstm_pytorch.blocks.slstm.block module xlstm_jax.models.xlstm_pytorch.blocks.slstm.cell module xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.cuda_init module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla.lstm module xlstm_jax.models.xlstm_pytorch.blocks.slstm.src.vanilla.slstm module xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block module xlstm_jax.models.xlstm_pytorch.components module xlstm_jax.models.xlstm_pytorch.components.conv module xlstm_jax.models.xlstm_pytorch.components.feedforward module xlstm_jax.models.xlstm_pytorch.components.init module xlstm_jax.models.xlstm_pytorch.components.linear_headwise module xlstm_jax.models.xlstm_pytorch.components.ln module xlstm_jax.models.xlstm_pytorch.components.util module xlstm_jax.models.xlstm_pytorch.utils module xlstm_jax.models.xlstm_pytorch.xlstm_block_stack module xlstm_jax.models.xlstm_pytorch.xlstm_lm_model module xlstm_jax.resume_training module xlstm_jax.start_training module xlstm_jax.train_init_fns module xlstm_jax.trainer module xlstm_jax.trainer.base module xlstm_jax.trainer.base.param_utils module xlstm_jax.trainer.base.trainer module xlstm_jax.trainer.callbacks module xlstm_jax.trainer.callbacks.callback module xlstm_jax.trainer.callbacks.checkpointing module xlstm_jax.trainer.callbacks.extended_evaluation module xlstm_jax.trainer.callbacks.lr_monitor module xlstm_jax.trainer.callbacks.profiler module xlstm_jax.trainer.data_module module xlstm_jax.trainer.llm module xlstm_jax.trainer.llm.sampling module xlstm_jax.trainer.llm.trainer module xlstm_jax.trainer.logger module xlstm_jax.trainer.logger.base_logger module xlstm_jax.trainer.logger.cmd_logging module xlstm_jax.trainer.logger.file_logger module xlstm_jax.trainer.logger.tensorboard_logger module xlstm_jax.trainer.logger.wandb_logger module xlstm_jax.trainer.metrics module xlstm_jax.trainer.optimizer module xlstm_jax.trainer.optimizer.ademamix module xlstm_jax.trainer.optimizer.optimizer module xlstm_jax.trainer.optimizer.scheduler module xlstm_jax.utils module xlstm_jax.utils.error_logging_utils module xlstm_jax.utils.model_param_handling module xlstm_jax.utils.model_param_handling.convert_checkpoint module xlstm_jax.utils.model_param_handling.convert_state_dict module xlstm_jax.utils.model_param_handling.handle_mlstm_simple module xlstm_jax.utils.model_param_handling.load module xlstm_jax.utils.model_param_handling.store module xlstm_jax.utils.pytree_utils module xlstm_norm (xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlock attribute) (xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlock attribute) (xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlock attribute) xLSTMBlock (class in xlstm_jax.models.xlstm_clean.blocks.xlstm_block) (class in xlstm_jax.models.xlstm_parallel.blocks.xlstm_block) (class in xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block) xLSTMBlockConfig (class in xlstm_jax.models.xlstm_clean.blocks.xlstm_block) (class in xlstm_jax.models.xlstm_parallel.blocks.xlstm_block) (class in xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block) xLSTMBlockStack (class in xlstm_jax.models.xlstm_clean.xlstm_block_stack) (class in xlstm_jax.models.xlstm_parallel.xlstm_block_stack) (class in xlstm_jax.models.xlstm_pytorch.xlstm_block_stack) xLSTMBlockStackConfig (class in xlstm_jax.models.xlstm_clean.xlstm_block_stack) (class in xlstm_jax.models.xlstm_parallel.xlstm_block_stack) (class in xlstm_jax.models.xlstm_pytorch.xlstm_block_stack) xLSTMLMModel (class in xlstm_jax.models.xlstm_clean.xlstm_lm_model) (class in xlstm_jax.models.xlstm_parallel.xlstm_lm_model) (class in xlstm_jax.models.xlstm_pytorch.xlstm_lm_model) xLSTMLMModelConfig (class in xlstm_jax.models.xlstm_clean.xlstm_lm_model) (class in xlstm_jax.models.xlstm_parallel.xlstm_lm_model) (class in xlstm_jax.models.xlstm_pytorch.xlstm_lm_model)

Z

zgate (xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayer attribute)