### Running a Specific Test for xlstm-jax (Bash)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/installation.md

This command executes a specific Pytest test file to verify the correct installation and basic functionality of the `xlstm-jax` project. It is used to confirm the setup works correctly on the CPU.

```bash
pytest tests/config/test_equivalence_hydra_nonhydra.py
```

--------------------------------

### Cloning the xlstm-jax Repository (Bash)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/installation.md

This snippet clones the `xlstm-jax` repository from GitHub and navigates into the newly created directory. It is the essential first step to obtain the project's source code locally.

```bash
git clone https://github.com/NX-AI/xlstm-jax.git
cd xlstm-jax
```

--------------------------------

### Installing xlstm-jax in Editable Mode (Bash)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/installation.md

This command installs the `xlstm-jax` package in editable mode (`-e`), allowing local changes to the source code to be reflected without reinstallation. It is typically run after setting up the environment.

```bash
pip install -e .
```

--------------------------------

### Creating Mamba Environment for xlstm-jax (Bash)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/installation.md

This command, similar to the Conda command, creates a new environment using Mamba, a faster alternative to Conda. It utilizes the same YAML file to efficiently install all project dependencies.

```bash
mamba env create -f envs/environment_python_3.11_jax_0.4.34_cuda_12.6.yml
```

--------------------------------

### Running a Specific Test on GPU for xlstm-jax (Bash)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/installation.md

This command runs the same Pytest test as the previous snippet but explicitly targets a specific GPU (device 0) using `CUDA_VISIBLE_DEVICES`. This verifies GPU-accelerated functionality and correct CUDA setup.

```bash
CUDA_VISIBLE_DEVICES=0 pytest tests/config/test_equivalence_hydra_nonhydra.py
```

--------------------------------

### Creating Conda Environment for xlstm-jax (Bash)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/installation.md

This command creates a new Conda environment using the specified YAML file, which contains all necessary Python and JAX dependencies. This ensures a consistent and isolated development environment for the project.

```bash
conda env create -f envs/environment_python_3.11_jax_0.4.34_cuda_12.6.yml
```

--------------------------------

### Training mLSTM-v1 Model with Hydra on SlimPajama-627B (Bash)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/example_training.md

This command initiates training for an mLSTM-v1 model with 165M parameters on the SlimPajama-627B dataset using Hydra for hyperparameter management. It leverages a predefined experiment configuration file, `train_mLSTMv1_165M_slimpajama627b.yaml`, which specifies data, model, parallelization, and optimizer settings. This is the recommended approach for managing complex training setups.

```bash
PYTHONPATH=. python scripts/training/train_with_hydra.py +experiment=train_mLSTMv1_165M_slimpajama627b
```

--------------------------------

### Generating Example Input for JAX xLSTM Model

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Generates a random integer array representing example input tokens for the xLSTM model. The array's shape is determined by the defined batch size and context length, with values within the vocabulary range.

```Python
exmp_input = jax.random.randint(jax.random.PRNGKey(0), (BATCH_SIZE, CONTEXT_LENGTH), minval=0, maxval=VOCAB_SIZE)
```

--------------------------------

### Saving JAX Model Outputs and Inputs to NPZ File

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Saves the retrieved JAX logits and example inputs into a compressed NumPy `.npz` file. This provides a convenient way to persist the model's outputs and corresponding inputs.

```Python
np.savez("./logits_inputs_jax.npz", logits_jax=logits_jax_np, inputs=example_inputs_np)
```

--------------------------------

### Retrieving JAX Model Outputs to NumPy

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Transfers the computed JAX logits and the example input array from JAX device memory to standard NumPy arrays. This allows for easier manipulation and saving of the results.

```Python
logits_jax_np = jax.device_get(logits_jax)
example_inputs_np = jax.device_get(exmp_input)
```

--------------------------------

### Defining Input Dimensions for JAX xLSTM Model

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Extracts the vocabulary size and context length from the loaded JAX configuration and defines a batch size. These dimensions are used to construct example input tensors for the model.

```Python
VOCAB_SIZE = jax_config.vocab_size
BATCH_SIZE = 3
CONTEXT_LENGTH = jax_config.context_length
VOCAB_SIZE, BATCH_SIZE, CONTEXT_LENGTH
```

--------------------------------

### Performing JAX xLSTM Model Forward Pass with Shard Map

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Sets up the model variables using the loaded checkpoint and defines a JAX-jitted forward pass function. This function applies the xLSTM model to the example input, utilizing `shard_map` for distributed inference.

```Python
variables = {}
variables["params"] = flax.core.frozen_dict.unfreeze(jax_checkpoint)

def _forward(
    batch_input: jax.Array, variables: Any, batch_position: jax.Array | None, batch_borders: jax.Array | None
) -> jax.Array:
    return xlstm_model_jax.apply(
        variables,
        batch_input,
        pos_idx=batch_position,
        document_borders=batch_borders,
        train=True,
        rngs={"dropout": jax.random.PRNGKey(42)},
    )


forward_fn = jax.jit(
    shard_map(
        _forward,
        mesh,
        in_specs=(P(), variables_partition_specs, P(), P()),
        out_specs=P(),
        check_rep=False,
    ),
)

logits_jax = forward_fn(exmp_input, variables, None, None)
```

--------------------------------

### Executing Hydra Experiment with SLURM Launcher

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This command demonstrates how to run a Hydra-based training script (`train_with_hydra.py`) on a SLURM cluster using the `submitit_launcher`. The `--multirun hydra/launcher=slurm_launcher` argument activates the SLURM integration, and `+experiment={YOUR_EXPERIMENT_FILE}` specifies the experiment configuration file to be used.

```Shell
PYTHONPATH=. python scripts/training/train_with_hydra.py --multirun hydra/launcher=slurm_launcher +experiment={YOUR_EXPERIMENT_FILE}
```

--------------------------------

### Generating Resume Command for xLSTM-JAX Experiment (Basic)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This command initiates the process of generating a CLI command to resume an xLSTM-JAX training experiment. It requires the path to the previous experiment's output folder to retrieve its configuration.

```Python
python scripts/traininig/get_cli_command_to_resume_training.py --resume_from_folder=PATH_TO_RUN
```

--------------------------------

### Generating Resume Command for xLSTM-JAX Experiment (Advanced with SLURM and Overrides)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This comprehensive command generates a CLI command to resume an xLSTM-JAX training experiment, incorporating SLURM for job management, specifying a particular checkpoint step, and applying new configuration overrides for training steps, learning rate, and logging frequency.

```Python
python scripts/training/get_cli_command_to_resume_training.py --resume_from_folder=PATH_TO_RUN --use_slurm --checkpoint_step=95000 --new_overrides="num_train_steps=20000 lr=0.0001 logger.log_every_n_steps=10"
```

--------------------------------

### Default SLURM Launcher Configuration for Hydra

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This YAML configuration defines the default parameters for Hydra's `submitit_slurm` launcher, specifying resource allocation (CPUs, GPUs, memory), job naming, partition, and various NCCL/UCX environment variables essential for distributed training on a SLURM cluster. It includes settings for GPU binding, job timeout, and specific network optimizations.

```yaml
# @package _global_
defaults:
  - override /hydra/launcher: submitit_slurm

hydra:
  launcher:
    submitit_folder: ${hydra.sweep.dir}/.submitit/%j
    timeout_min: 60
    cpus_per_task: 28
    gpus_per_task: 1   # each GPU must have its separate task for jax
    gres: gpu:${n_gpus}
    tasks_per_node: ${n_gpus}
    mem_gb: ~
    nodes: 1
    name: ${hydra.job.name}
    _target_: hydra_plugins.hydra_submitit_launcher.submitit_launcher.SlurmLauncher
    partition: compute
    qos: null
    comment: testing_slurmit_launcher
    additional_parameters: {
      "gpu-bind": "closest",
      "wait-all-nodes": "1",
      "time": "7-00:00:00",
      "exclusive": "",
    }
    srun_args:
      - "--kill-on-bad-exit=1"
    setup:
      - export NCCL_CROSS_NIC=0
      - export NCCL_SOCKET_NTHREADS=16
      - export NCCL_DEBUG=WARN
      - export NCCL_CUMEM_ENABLE=0
      - export NCCL_IB_SPLIT_DATA_ON_QPS=0
      - export NCCL_IB_QPS_PER_CONNECTION=16
      - export NCCL_IB_GID_INDEX=3
      - export NCCL_IB_TC=41
      - export NCCL_IB_SL=0
      - export NCCL_IB_TIMEOUT=22
      - export NCCL_NET_PLUGIN=none
      - export NCCL_SOCKET_IFNAME=eth0
      - export NCCL_IGNORE_CPU_AFFINITY=1
      - export NCCL_IB_HCA="=mlx5_0,mlx5_1,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7,mlx5_8,mlx5_9,mlx5_10,mlx5_12,mlx5_13,mlx5_14,mlx5_15,mlx5_16,mlx5_17"
      - export HCOLL_ENABLE_MCAST_ALL=0
      - export coll_hcoll_enable=0
      - export UCX_TLS=tcp
      - export UCX_NET_DEVICES=eth0
      - export RX_QUEUE_LEN=8192
      - export IB_RX_QUEUE_LEN=8192
      - export OMPI_MCA_coll=^hcoll
      - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/nfs-gpu/xlstm/miniforge3/envs/python_3.11_jax_0.4.34_cuda_12.6/cuda-compat
      - export TOKENIZERS_PARALLELISM=0
      - export JAX_COMPILATION_CACHE_DIR=/nfs-gpu/xlstm/data/jax_compilation_cache
```

--------------------------------

### Importing Core Libraries and mLSTM Components

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet imports essential libraries like `pathlib`, `numpy`, and `torch`, along with specific components from the `mlstm_simple` package, including `load_from_pretrained` for model loading and `mLSTM`/`mLSTMConfig` for model definition.

```python
from pathlib import Path

import numpy as np
import torch
from mlstm_simple.from_pretrained import load_from_pretrained
from mlstm_simple.model import mLSTM, mLSTMConfig
```

--------------------------------

### Running Hydra Experiment File from Command Line in Python

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This command illustrates how to execute a Hydra-enabled Python script using a predefined experiment configuration file. The `+experiment` syntax tells Hydra to load and apply the specified experiment file, enabling reproducible experiment runs.

```bash
PYTHONPATH=. python scripts/training/train_with_hydra.py +experiment=synthetic_experiment
```

--------------------------------

### Instantiating JAX xLSTM Language Model

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Creates an instance of the `xLSTMLMModel` using the loaded JAX model configuration. This prepares the model architecture for parameter initialization and subsequent inference.

```Python
xlstm_model_jax = xLSTMLMModel(jax_config)
```

--------------------------------

### Initializing Hydra Main Function in Python

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This Python decorator marks the `main_train` function as the entry point for Hydra, specifying the configuration path and the default configuration file name. It ensures Hydra loads and processes the configuration before the function execution.

```python
@hydra.main(config_path="../configs", config_name="config", version_base="1.3")
def main_train(cfg: DictConfig):
    ...
```

--------------------------------

### Defining Base Configuration with Hydra Defaults in YAML

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This YAML snippet illustrates a base Hydra configuration file (`config.yaml`) defining default values and referencing other configuration schemas and groups. It specifies the initial device setting and how other config modules (like `parallel` and `model`) are integrated.

```yaml
defaults:
  - config_schema
  - parallel: synthetic
  - model: mLSTM120M
  - _self_


# General hyperparameters. Will be put in their respective config modules once they are created.

# device. cpu or gpu
device: cpu
```

--------------------------------

### Training LLama Baseline Model without Hydra on SlimPajama (Bash)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/example_training.md

This command trains a LLama baseline model on the SlimPajama dataset without Hydra. Users must specify a log directory and choose a model configuration, typically '1.3B' or '165M', as defined within the script. This method provides a direct way to initiate training for LLama models.

```bash
PYTHONPATH=. python scripts/training/run_train_llama_slimpajama.py --log_dir=<log_dir> --model=<model>
```

--------------------------------

### Checking Hydra Configuration Types in Python

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This command demonstrates how to use a utility script to compile and log a Hydra configuration, primarily for type checking. It helps validate experiment files before running resource-intensive jobs, ensuring parameters are correctly typed.

```bash
python scripts/check_config.py +experiment={YOUR_EXPERIMENT_FILE}
```

--------------------------------

### Loading JAX Logits and Inputs from NPZ

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet loads a NumPy archive file (`.npz`) containing pre-computed JAX logits and input data, which will be used later for comparison and model inference.

```python
logits_inputs_jax = np.load("./logits_inputs_jax.npz")
```

--------------------------------

### Pretty Printing mLSTM Model Configuration

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet imports the `pprint` module and uses it to display the configuration object (`model.config`) of the loaded mLSTM model in a human-readable, formatted way, revealing its various parameters.

```python
from pprint import pprint

pprint(model.config)
```

--------------------------------

### Displaying Flattened JAX Model Configuration

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Prints the loaded JAX model configuration in a flattened dictionary format using `pprint` for enhanced readability. This helps in inspecting all configuration parameters.

```Python
from pprint import pprint

pprint(flatten_dict(asdict(jax_config)))
```

--------------------------------

### Configuring Python Path for mLSTM Imports

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet modifies the Python system path to include necessary directories for importing `mlstm_simple` and related modules, ensuring that local or project-specific packages can be found.

```python
import sys

sys.path.append("../..")
sys.path.append("../../../mlstm_simple_torch")
```

--------------------------------

### Defining Parallelism Configuration with Python Dataclass

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This Python dataclass, `ParallelConfig`, defines various parameters for configuring parallelism, including sizes and names for data, FSDP, pipeline, and model axes. It also specifies settings for activation checkpointing, FSDP sharding, and parameter gathering/scattering data types. A `__post_init__` method ensures validation of the FSDP dtype settings.

```Python
@dataclass(kw_only=True, frozen=False)
class ParallelConfig:
    """Configuration for parallelism."""

    data_axis_size: int = -1
    """Size of the data axis. If -1, it will be inferred by the number of available devices."""
    fsdp_axis_size: int = 1
    """Size of the FSDP axis. If -1, it will be inferred by the number of available devices."""
    pipeline_axis_size: int = 1
    """Size of the pipeline axis. If -1, it will be inferred by the number of available devices."""
    model_axis_size: int = 1
    """Size of the model axis. If -1, it will be inferred by the number of available devices."""
    data_axis_name: str = "dp"
    """Name of the data axis."""
    fsdp_axis_name: str = "fsdp"
    """Name of the FSDP axis."""
    pipeline_axis_name: str = "pp"
    """Name of the pipeline axis."""
    model_axis_name: str = "tp"
    """Name of the model axis."""
    remat: list[str] = field(default_factory=lambda: [])
    """Module names on which we apply activation checkpointing / rematerialization."""
    fsdp_modules: list[str] = field(default_factory=lambda: [])
    """Module names on which we apply FSDP sharding."""
    fsdp_min_weight_size: int = 2**18
    """Minimum size of a parameter to be sharded with FSDP."""
    fsdp_gather_dtype: str | None = None
    """The dtype to cast the parameters to before gathering with FSDP. If `None`, no casting is performed and parameters
    are gathered in original precision (for example `float32`)."""
    fsdp_grad_scatter_dtype: str | None = None
    """The dtype to cast the gradients to before scattering. If `None`, the dtype of the parameters is used."""
    tp_async_dense: bool = False
    """Whether to use asynchronous tensor parallelism for dense layers. Default to `False`, as on local hardware,
    ppermute communication introduces large overhead."""

    def __post_init__(self):
        _allowed_fsdp_dtypes = ["float32", "bfloat16", "float16"]

        if self.fsdp_gather_dtype is not None:
            assert self.fsdp_gather_dtype in _allowed_fsdp_dtypes
        if self.fsdp_grad_scatter_dtype is not None:
            assert self.fsdp_grad_scatter_dtype in _allowed_fsdp_dtypes

```

--------------------------------

### Initializing JAX Device Mesh for Parallelism

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Initializes a JAX device mesh based on the defined `ParallelConfig` and a subset of available JAX devices. This mesh is fundamental for orchestrating distributed computations.

```Python
mesh = initialize_mesh(parallel_config=parallel, device_array=np.array(jax.devices())[0:1])
```

--------------------------------

### Enabling SLURM Launcher in Hydra Experiment File

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This YAML snippet shows how to explicitly configure Hydra to use the `submitit_launcher` directly within an experiment configuration file. By overriding `/hydra/launcher` to `slurm_launcher` and setting `hydra.mode` to `MULTIRUN`, the experiment will automatically be submitted to SLURM without requiring command-line arguments.

```yaml
  - override /hydra/launcher: slurm_launcher

hydra:
  mode: MULTIRUN
```

--------------------------------

### Training xLSTM Model without Hydra on SlimPajama (Bash)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/example_training.md

This command trains an xLSTM model on the SlimPajama dataset without using Hydra. It requires specifying a log directory for checkpoints and a predefined model configuration. The `--use_full_dataset` flag can be added to train on the larger SlimPajama-627B dataset instead of the default subset.

```bash
PYTHONPATH=. python scripts/training/run_train_slimpajama.py --log_dir=<log_dir> --model=<model>
```

--------------------------------

### Generating JAX to PyTorch Checkpoint Conversion Command

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet constructs a shell command string to execute a Python script responsible for converting a JAX model checkpoint to a PyTorch format. It dynamically inserts the previously defined JAX checkpoint path and the desired PyTorch output path into the command.

```python
# ## Convert jax checkpoint to torch:
command = f'PYTHONPATH=. python scripts/checkpoint_conversion/convert_mlstm_checkpoint_jax_to_torch_simple.py --checkpoint_dir "{str(JAX_CHECKPOINT_PATH)}" --output_path "{str(SAVE_TORCH_CHECKPOINT_AT)}" --checkpoint_type plain'
print(command)
```

--------------------------------

### Defining JAX Checkpoint Path

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet defines a constant string `JAX_CHECKPOINT_PATH` that specifies the absolute file system path to a JAX model checkpoint, which will be used as the source for conversion to a PyTorch model.

```python
JAX_CHECKPOINT_PATH = "/nfs-gpu/xlstm/logs/outputs/xlstm-jax/DCLM/dclm_mLSTMv1_1.3B_ctx8192_2024-11-19T09:24:50/0/checkpoints/checkpoint_95000"
```

--------------------------------

### Summarizing Argmax Prediction Agreement

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet calculates and displays the count of matching argmax predictions and the total number of predictions for the first batch, providing a quantitative measure of agreement between the PyTorch and JAX models' top predictions.

```python
batch1_equal_argmax.sum(), len(batch1_equal_argmax)
```

--------------------------------

### Importing xLSTM JAX Components and Utilities

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Imports essential modules and classes required for defining, configuring, and interacting with the xLSTM model in JAX, including distributed utilities, model configurations, and checkpoint handling functions.

```Python
from dataclasses import asdict
from typing import Any

import flax
import jax
import numpy as np
from flax import linen as nn
from jax.experimental.shard_map import shard_map
from jax.sharding import PartitionSpec as P

from xlstm_jax.distributed.mesh_utils import initialize_mesh
from xlstm_jax.models.configs import ParallelConfig
from xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend import mLSTMBackendNameAndKwargs
from xlstm_jax.models.xlstm_parallel.blocks.mlstm.block import mLSTMBlockConfig
from xlstm_jax.models.xlstm_parallel.blocks.mlstm.cell import mLSTMCellConfig
from xlstm_jax.models.xlstm_parallel.blocks.mlstm.layer import mLSTMLayerConfig
from xlstm_jax.models.xlstm_parallel.components.feedforward import FeedForwardConfig
from xlstm_jax.models.xlstm_parallel.xlstm_lm_model import xLSTMLMModel, xLSTMLMModelConfig
from xlstm_jax.utils.model_param_handling.convert_checkpoint import convert_orbax_checkpoint_to_torch_state_dict
from xlstm_jax.utils.model_param_handling.handle_mlstm_simple import (
    pipeline_convert_mlstm_checkpoint_jax_to_torch_simple,
)
from xlstm_jax.utils.model_param_handling.load import load_model_params_and_config_from_checkpoint
from xlstm_jax.utils.pytree_utils import flatten_dict
```

--------------------------------

### Initializing JAX xLSTM Model Parameters with Shard Map

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Defines and executes a JAX-jitted function to initialize the xLSTM model's parameters. It leverages `shard_map` to distribute parameter initialization across the device mesh, ensuring proper sharding.

```Python
def _init_model(rng: jax.Array, batch_input: jax.Array) -> Any:
    param_rng, dropout_rng = jax.random.split(rng)
    # Initialize parameters.
    variables = xlstm_model_jax.init({"params": param_rng, "dropout": dropout_rng}, batch_input)
    return variables


# Prepare PRNG.
init_rng = jax.random.PRNGKey(42)
# First infer the output sharding to set up shard_map correctly.
# This does not actually run the init, only evaluates the shapes.
init_model_fn = jax.jit(
    shard_map(
        _init_model,
        mesh,
        in_specs=(P(), P()),
        out_specs=P(),
        check_rep=False,
    ),
)
variables_shapes = jax.eval_shape(init_model_fn, init_rng, exmp_input)
variables_partition_specs = nn.get_partition_spec(variables_shapes)
# Run init model function again with correct output specs.
init_model_fn = jax.jit(
    shard_map(
        _init_model,
        mesh,
        in_specs=(P(), P()),
        out_specs=variables_partition_specs,
        check_rep=False,
    ),
)

variables = init_model_fn(init_rng, exmp_input)
```

--------------------------------

### Running PyTorch Model Inference with Optimization

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet conditionally compiles the PyTorch model using `torch.compile` for performance. It then performs a forward pass on the model using the JAX inputs (converted to a PyTorch tensor and moved to CUDA), optionally within an `autocast` context for mixed precision, and captures the resulting logits and states.

```python
if USE_TORCH_COMPILE:
    model = torch.compile(model)
with torch.autocast(device_type="cuda", dtype=TORCH_AMP_DTYPE, enabled=ENABLE_TORCH_AMP):
    logits_torch, state = model(torch.from_numpy(inputs_jax).to("cuda"))
```

--------------------------------

### Displaying PyTorch Model Object

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet simply outputs the `model` object, which typically triggers its `__repr__` method in an interactive environment, displaying a summary of the loaded PyTorch mLSTM model's architecture and configuration.

```python
model
```

--------------------------------

### Overriding Hydra Configuration from Command Line in Python

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This command demonstrates how to execute a Hydra-enabled Python script while overriding specific configuration parameters directly from the command line. It allows for dynamic adjustment of settings like device and model parameters without modifying configuration files.

```bash
PYTHONPATH=. python scripts/training/train_with_hydra.py device=gpu model.num_heads=42
```

--------------------------------

### Defining PyTorch Checkpoint Save Path

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet constructs and resolves the file path where the converted PyTorch model checkpoint will be saved. It uses `pathlib.Path` to create a platform-independent path relative to the current directory.

```python
SAVE_TORCH_CHECKPOINT_AT = (Path(".").parent / "mlstm_simple_checkpoint").resolve()
SAVE_TORCH_CHECKPOINT_AT
```

--------------------------------

### Configuring Python Path for xLSTM JAX

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Adds the parent directory to the Python system path, enabling the import of local modules and packages from the xlstm_jax project structure.

```Python
import sys

sys.path.append("../../..")
```

--------------------------------

### Extracting First Batch Logits for Comparison

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet extracts the logits for the first batch (index 0) from both the PyTorch-derived NumPy array (`logits_torch_np`) and the original JAX logits (`logits_jax`), preparing them for element-wise comparison.

```python
torch_logits_batch1 = logits_torch_np[0]
jax_logits_batch1 = logits_jax[0]
```

--------------------------------

### Citing the xLSTM-jax Repository (BibTeX)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/README.md

This BibTeX entry is for citing the xLSTM-jax GitHub repository itself. It includes the title, author (NXAI GmbH), publication year, and the URL to the repository, suitable for referencing the codebase.

```BibTeX
@misc{xlstm-jax,
  title={xLSTM-jax},
  author={NXAI GmbH},
  year={2024},
  url={https://github.com/NX-AI/xlstm-jax/},
}
```

--------------------------------

### Loading JAX Model Checkpoint and Configuration

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Loads the model parameters (checkpoint) and its associated configuration from the specified JAX checkpoint path. The configuration is returned as a dataclass for structured access.

```Python
jax_checkpoint, jax_config = load_model_params_and_config_from_checkpoint(
    JAX_CHECKPOINT_PATH, return_config_as_dataclass=True
)
```

--------------------------------

### Configuring PyTorch Performance Settings

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet sets various PyTorch configuration options, including the data type for Automatic Mixed Precision (AMP), whether to enable AMP, and whether to use `torch.compile` for performance optimization. It also sets the float32 matrix multiplication precision to 'high'.

```python
TORCH_AMP_DTYPE = torch.float32
ENABLE_TORCH_AMP = False
USE_TORCH_COMPILE = True
torch.set_float32_matmul_precision(
    "high"
)  # TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance
```

--------------------------------

### Loading PyTorch mLSTM Model from Pretrained Checkpoint

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet loads an mLSTM model into PyTorch from the previously converted checkpoint. It specifies various kernel names and a chunk size, configuring the model for specific inference optimizations.

```python
model = load_from_pretrained(
    checkpoint_path=SAVE_TORCH_CHECKPOINT_AT,
    chunkwise_kernel_name="chunkwise--triton_xl_chunk",
    sequence_kernel_name="native_sequence__triton_step_fused",
    step_kernel_name="triton_fused",
    chunk_size=128,
)
```

--------------------------------

### Converting Hugging Face Dataset to ArrayRecord (Python)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/dataset_preparation.md

This command executes a Python script to convert a specified Hugging Face dataset, such as 'DKYoon/SlimPajama-6B', into the ArrayRecord format. This conversion is a crucial step in preparing data for the xLSTM-JAX project's preprocessing pipeline. For standard pre-training datasets, the script primarily reads and converts the text column.

```Shell
PYTHONPATH=. python scripts/data_processing/hf_to_arrayrecord.py --hf_path=DKYoon/SlimPajama-6B
```

--------------------------------

### Registering ParallelConfig with Hydra's Config Store

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This Python snippet demonstrates how to register the `ParallelConfig` dataclass with Hydra's config store. By mapping the string 'parallel_schema' to the `ParallelConfig` class, Hydra can automatically use its defined attributes and default values when this schema is referenced in configuration YAML files.

```Python
cs.store(name="parallel_schema", group="parallel", node=ParallelConfig)
```

--------------------------------

### Extracting Top 5 Predictions for PyTorch and JAX

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet identifies the indices of the top 5 predicted classes (highest logits) for each token in the first batch, separately for both the PyTorch and JAX models. This allows for a more detailed comparison beyond just the single top prediction.

```python

```

--------------------------------

### Displaying JAX Logits Array

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet simply outputs the `logits_jax` NumPy array, which contains the original logits computed by the JAX model, allowing for direct inspection or comparison in an interactive environment.

```python
logits_jax
```

--------------------------------

### Extracting JAX Logits and Inputs

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet extracts the 'logits_jax' and 'inputs' arrays from the previously loaded `logits_inputs_jax` NumPy archive, making them accessible for further processing.

```python
logits_jax = logits_inputs_jax["logits_jax"]
inputs_jax = logits_inputs_jax["inputs"]
```

--------------------------------

### Splitting ArrayRecord Dataset for Training/Validation (Python)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/dataset_preparation.md

This command runs a Python script to split an existing ArrayRecord dataset, specifically the DCLM dataset, into distinct training and validation sets. The script generates a 500k sequence validation split from the original single-split dataset, producing new ArrayRecord files for both partitions.

```Shell
PYTHONPATH=. python scripts/data_processing/split_array_records_dataset.py --dataset_name=DCLM
```

--------------------------------

### Defining JAX Checkpoint Path

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Specifies the absolute file system path to a pre-trained JAX xLSTM model checkpoint. This path is crucial for loading the model's saved parameters and configuration.

```Python
JAX_CHECKPOINT_PATH = "/nfs-gpu/xlstm/logs/outputs/xlstm-jax/DCLM/dclm_mLSTMv1_1.3B_ctx8192_2024-11-19T09:24:50/0/checkpoints/checkpoint_95000"
```

--------------------------------

### Displaying Argmax Equality Comparison Result

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet outputs the boolean array `batch1_equal_argmax`, which shows element-wise whether the top predicted class from the PyTorch model matches that of the JAX model for the first batch.

```python
batch1_equal_argmax
```

--------------------------------

### Moving Model to CUDA and Configuring State Return

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet moves the loaded PyTorch model to the GPU ('cuda') for accelerated computation. It also sets a configuration option to ensure the model returns its last hidden states during inference.

```python
model = model.to("cuda")
model.config.return_last_states = True
```

--------------------------------

### Defining Hydra Experiment Overrides in YAML

Source: https://github.com/nx-ai/xlstm-jax/blob/main/docs/configuration_with_hydra.md

This YAML snippet shows an experiment-specific configuration file that overrides global defaults and sets specific parameters for a reproducible experiment. It uses `override` keywords to ensure certain config groups are replaced and defines task-specific hyperparameters.

```yaml
# @package _global_
defaults:
  - /data@data_train.ds1: synthetic
  - override /parallel: synthetic
  - override /model: mLSTMv1_165M
  - _self_

# specify the deltas from the defaults:
task_name: slurm_tests
batch_size_per_device: 2
context_length: 128
num_train_steps: 10
lr: 1e-3

logger:
  log_every_n_steps: 2
```

--------------------------------

### Configuring JAX Parallelism for xLSTM Model

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Defines the `ParallelConfig` for the JAX xLSTM model, specifying axis names for data, FSDP, model, and pipeline parallelism. It also sets FSDP-specific parameters like gather dtype and minimum weight size.

```Python
parallel = ParallelConfig(
    data_axis_name="dp",
    fsdp_axis_name="fsdp",
    model_axis_name="tp",
    pipeline_axis_name="pp",
    fsdp_modules=[],
    fsdp_gather_dtype="bfloat16",
    fsdp_min_weight_size=2**18,
    remat=[],
    fsdp_axis_size=1,
    model_axis_size=1,
    data_axis_size=1,
    tp_async_dense=False,
)
```

--------------------------------

### Calculating Argmax for JAX Logits

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet computes the index of the maximum value along the last axis for the first batch of JAX logits, effectively determining the predicted class for each token in the sequence from the original JAX model.

```python
np.argmax(jax_logits_batch1, axis=-1)
```

--------------------------------

### Comparing Argmax Predictions Between PyTorch and JAX

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet performs a boolean comparison of the argmax predictions (most likely class indices) between the first batch of PyTorch logits and JAX logits. The result is a boolean array indicating where the predictions match.

```python
batch1_equal_argmax = np.argmax(torch_logits_batch1, axis=-1) == np.argmax(jax_logits_batch1, axis=-1)
```

--------------------------------

### Loading Data from NPZ File

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Loads the content of the previously saved `.npz` file into a NumPy file object. This allows access to the individual arrays stored within the archive.

```Python
file = np.load("./logits_inputs_jax.npz")
```

--------------------------------

### Accessing Logits from Loaded NPZ File

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_jax_model.ipynb

Retrieves the 'logits_jax' array from the loaded `.npz` file object. This demonstrates how to access specific data components stored within the NumPy archive.

```Python
file["logits_jax"]
```

--------------------------------

### Calculating Argmax for PyTorch Logits

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet computes the index of the maximum value along the last axis for the first batch of PyTorch logits, effectively determining the predicted class for each token in the sequence.

```python
np.argmax(torch_logits_batch1, axis=-1)
```

--------------------------------

### Converting PyTorch Logits to NumPy Array

Source: https://github.com/nx-ai/xlstm-jax/blob/main/tests/models/mlstm_simple/load_pretrained_torch_model.ipynb

This snippet converts the PyTorch logits tensor, obtained from model inference, into a NumPy array. It first casts the tensor to `float`, moves it to the CPU, detaches it from the computation graph, and then converts it to a NumPy array for further analysis and comparison.

```python
logits_torch_np = logits_torch.float().cpu().detach().numpy()
```

--------------------------------

### Citing the xLSTM Paper (BibTeX)

Source: https://github.com/nx-ai/xlstm-jax/blob/main/README.md

This BibTeX entry provides the necessary information to cite the xLSTM research paper, including authors, title, journal, and publication year. It is used for academic referencing when utilizing the xLSTM model or concepts.

```BibTeX
@article{xlstm,
  title={xLSTM: Extended Long Short-Term Memory},
  author={Beck, Maximilian and P{\"o}ppel, Korbinian and Spanring, Markus and Auer, Andreas and Prudnikova, Oleksandra and Kopp, Michael and Klambauer, G{\"u}nter and Brandstetter, Johannes and Hochreiter, Sepp},
  journal={arXiv preprint arXiv:2405.04517},
  year={2024}
}
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.