### Pretrain Neuro-JEPA (First 200 Epochs)

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Starts the pretraining process for the first 200 epochs with an annealing schedule. Assumes a single node and GPU setup.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
  scripts/pretrain.py --config-name pretrain_neurojepa_base
```

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

### Install with Python Virtual Environment

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Use this method to install the project dependencies within a dedicated Python virtual environment. Ensure Python 3.11+ is installed.

```bash
python -m venv .venv
source .venv/bin/activate
pip install -U pip
pip install -r requirements.txt
pip install -e .
```

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

### Install with Conda Environment

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

This method installs the project dependencies using a Conda environment, ensuring Python 3.11 is used.

```bash
conda create -n neurojepa python=3.11 -y
conda activate neurojepa
pip install -U pip
pip install -r requirements.txt
pip install -e .
```

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

### Pretrain Neuro-JEPA (Cooldown Phase)

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Initiates the cooldown phase of pretraining for 40 epochs. Assumes a single node and GPU setup.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
  scripts/pretrain.py --config-name cooldown_neurojepa_base
```

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

### Create DataLoader and Extract Batch

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Initializes a MONAI Dataset and DataLoader for batch processing. It then iterates through the DataLoader to get a batch of preprocessed images and moves them to the appropriate device (e.g., CUDA).

```python
# Define dataloader and get scan array
batch_size = 2

test_ds = data.Dataset(
    data=sample_images, 
    transform=trans,
)

test_loader = data.DataLoader(
    dataset=test_ds,
    batch_size=batch_size,
    num_workers=1,
    pin_memory=True,
    shuffle=False,
)

x_img = next(iter(test_loader))['image']
x_img = x_img.to(device)
```

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

### Extract Feature Example Notebook

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Jupyter notebook demonstrating how to load a trained Neuro-JEPA model and extract features from MRI scans. This is useful for downstream tasks.

```python
# Example usage: Open and run the cells in ./notebooks/extract_feature_example.ipynb

```

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

### Load Neuro-JEPA Backbone

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Loads the Neuro-JEPA backbone model from Hugging Face. Ensure you have the 'neurojepa' library installed and a CUDA-enabled GPU for optimal performance.

```python
import torch
from neurojepa.utils.init_utils import load_backbone_from_hf

device = "cuda" if torch.cuda.is_available() else "cpu"

backbone = load_backbone_from_hf(
    "NYUMedML/Neuro-JEPA",
    device=device,
)

backbone.eval()
```

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

### Inference Script

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Python script for running inference with a trained Neuro-JEPA model. This script can be used to get predictions on new data.

```python
# Example usage: python ./scripts/inference.py --config-path ./configs/inference/ --config-name <your_config_name>

```

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

### Pretrain with Explicit Data and Checkpoint Overrides

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Runs the pretraining script with custom paths for training data, cache directory, and checkpoint saving. Useful for specifying custom datasets or storage locations.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
  --config-name pretrain_neurojepa_base \
  data.train_csv_path=/path/to/pretrain.csv \
  data.cache_dir=/path/to/cache \
  meta.save_path_checkpoint=/path/to/checkpoints/neurojepa
```

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

### Monitor Neuro-JEPA Model Performance

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

This command initiates the monitoring procedure for Neuro-JEPA model performance across multiple tasks using attentive probing. It requires a specific monitoring configuration file.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
  scripts/monitor.py --config-name monitor_t1w_nyu.yaml
```

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

### Define Sample Image Data

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Creates a list of dictionaries, where each dictionary represents a sample image with its file path. This data structure is used to initialize the MONAI Dataset.

```python
sample_images = [
    {'image': '/path/to/filename1.nii.gz'},
    {'image': '/path/to/filename2.nii.gz'},
]
```

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

### Load Backbone from Hugging Face

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Loads pre-trained model weights directly from Hugging Face. Ensure PyTorch and CUDA are available for GPU acceleration.

```python
import torch
from neurojepa.utils.init_utils import load_backbone_from_hf

device = "cuda" if torch.cuda.is_available() else "cpu"

backbone = load_backbone_from_hf(
    "NYUMedML/Neuro-JEPA",
    device=device,
)
```

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

### Run Inference with Hugging Face Checkpoint

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Executes the inference script loading weights directly from Hugging Face. Requires model ID and optionally a token.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
  scripts/inference.py --config-name inference_neurojepa \
  meta.load_checkpoint=true \
  meta.checkpoint_source=hf \
  meta.hf_model_id=NYUMedML/Neuro-JEPA \
  meta.hf_checkpoint_filename=model.safetensors \
  meta.hf_token=true
```

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

### Activate Python Virtual Environment on Windows

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Command to activate the Python virtual environment on Windows systems.

```powershell
.venv\Scripts\activate
```

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

### Run Inference with Local Checkpoint

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Executes the inference script using a local checkpoint. Specify the path to your checkpoint file.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
  scripts/inference.py --config-name inference_neurojepa \
  meta.load_checkpoint=true \
  meta.checkpoint_source=local \
  meta.load_path_checkpoint=/path/to/checkpoints/latest.pt
```

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

### Pretraining Script

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Python script for pretraining the Neuro-JEPA model. This is the main entrypoint for the large-scale MRI pretraining task.

```python
# Example usage: python ./scripts/pretrain.py --config-path ./configs/pretrain/ --config-name <your_config_name>

```

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

### Monitoring Script

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Python script for monitoring training or evaluation processes. This script can be used to track metrics and visualize progress.

```python
# Example usage: python ./scripts/monitor.py --config-path ./configs/monitor/ --config-name <your_config_name>

```

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

### Multimodal Fine-tuning with Product-of-Experts

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Python script for multimodal classification fine-tuning using the Product-of-Experts (PoE) approach. This method is used for fusing multimodal information.

```python
# Example usage: python ./scripts/finetune/mm_poe.py --config-path ./configs/finetune/ --config-name <your_config_name>

```

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

### Fine-Tune Neuro-JEPA with Product-of-Experts Fusion

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Fine-tune the Neuro-JEPA model for multimodal classification using Product-of-Experts (PoE) fusion. This command specifies the multimodal script and the relevant configuration.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
 scripts/finetune/mm_poe.py --config-name finetune_neurojepa_mm
```

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

### Fine-Tune Neuro-JEPA for Unimodal Classification/Regression

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Use this command to fine-tune the Neuro-JEPA model for unimodal classification or regression tasks. Ensure the correct configuration file is specified.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
  scripts/finetune/default.py --config-name finetune_neurojepa
```

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

### Multimodal Fine-tuning Script

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Python script for multimodal classification fine-tuning. This script handles combining information from multiple modalities.

```python
# Example usage: python ./scripts/finetune/mm.py --config-path ./configs/finetune/ --config-name <your_config_name>

```

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

### Submit Neuro-JEPA Pretraining Job with SLURM

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

This command submits a SLURM job for multinode distributed pretraining of the Neuro-JEPA model. Ensure to adjust account, partition, node count, and path configurations before submission.

```bash
sbatch scripts/slurm/pretrain_neurojepa.sh
```

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

### Unimodal Fine-tuning Script

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Python script for unimodal classification or regression fine-tuning. This is an entrypoint for training.

```python
# Example usage: python ./scripts/finetune/default.py --config-path ./configs/finetune/ --config-name <your_config_name>

```

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

### Import MONAI Data and Transforms

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Imports necessary modules from the MONAI library for data handling and transformations. This includes data loading and custom lambda transformations.

```python
import numpy as np

from monai import data
from monai import transforms
from monai.transforms import Lambdad
```

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

### Fine-Tune Neuro-JEPA for Time-to-Event Modeling

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Use this command to fine-tune the Neuro-JEPA model for time-to-event or survival modeling tasks. It points to the dedicated TTE script and configuration.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
 scripts/finetune/tte.py --config-name finetune_neurojepa_tte
```

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

### Time-to-Event Fine-tuning Script

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Python script for time-to-event or survival modeling fine-tuning. This script is designed for survival analysis tasks.

```python
# Example usage: python ./scripts/finetune/tte.py --config-path ./configs/finetune/ --config-name <your_config_name>

```

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

### Submit MRI Registration Jobs

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Use this command to submit the T1w MRI registration script as a SLURM array job. Ensure CSV_FILE, OUTPUT_DIR, and template paths are correctly configured in the script.

```bash
sbatch --array=0-9 registration/register_t1w.sh
```

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

### Fine-Tune Neuro-JEPA for Multimodal Classification

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

This command is for fine-tuning the Neuro-JEPA model on multimodal classification tasks. It utilizes a specific script for multimodal processing.

```bash
python -m torch.distributed.run --nnodes 1 --nproc_per_node 1 --master_port 12345 \
  scripts/finetune/mm.py --config-name finetune_neurojepa_mm
```

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

### Define Preprocessing Pipeline

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Defines a comprehensive image preprocessing pipeline using MONAI transforms. This pipeline handles loading, channel ensuring, NaN removal, orientation, spacing, cropping, resizing, intensity scaling, and casting to float32.

```python
# ---- Load and preprocess the volume (same pipeline as training) ----
def remove_nan(img):
    img[np.isnan(img)] = 0.0
    return img

roi_size = [96, 108, 96]

trans = transforms.Compose([
    transforms.LoadImaged(keys=['image'], image_only=False),
    transforms.EnsureChannelFirstd(keys=['image']),
    Lambdad(('image',), remove_nan),
    transforms.Orientationd(keys=['image'], axcodes='RAS'),
    transforms.Spacingd(keys=['image'], pixdim=(1.0, 1.0, 1.0), mode=[5]),
    transforms.CropForegroundd(
        keys=['image'], source_key='image',
        select_fn=lambda x: x > 0.0, margin=4, allow_smaller=True),
    transforms.ResizeWithPadOrCropd(
        keys=['image'], spatial_size=[180, 216, 180], mode='edge'),
    transforms.Resized(keys=['image'], spatial_size=[100, 120, 100]),
    transforms.CenterSpatialCropd(
        keys=["image"],
        roi_size=roi_size,
        allow_missing_keys=True,
    ),
    transforms.ScaleIntensityRangePercentilesd(
        keys=['image'], lower=0.5, upper=99.5, b_min=0, b_max=1, clip=True),
    transforms.CastToTyped(
        keys=["image"],
        dtype=np.float32,
        allow_missing_keys=True,
    ),
])
```

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

### MLP Layer in Neuro-JEPA

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Illustrates the implementation of a standard Multi-Layer Perceptron (MLP) layer used within the Neuro-JEPA model. This layer consists of two linear transformations with a GELU activation and dropout.

```python
MLP(
            (fc1): Linear(in_features=768, out_features=384, bias=True)
            (act): GELU(approximate='none')
            (fc2): Linear(in_features=384, out_features=768, bias=True)
            (drop): Dropout(p=0.0, inplace=False)
          )
```

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

### MRI Registration Scripts

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Shell scripts for registering T1w, T2w, and FLAIR MRI scans to MNI152 standard space. These scripts are SLURM-ready.

```shell
#!/bin/bash
# Example usage: bash ./registration/register_t1w.sh <input_nii_gz_file> <output_dir>

```

```shell
#!/bin/bash
# Example usage: bash ./registration/register_t2w.sh <input_nii_gz_file> <output_dir>

```

```shell
#!/bin/bash
# Example usage: bash ./registration/register_flair.sh <input_nii_gz_file> <output_dir>

```

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

### SLURM Pretraining Job Script

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Shell script for submitting a pretraining job to a SLURM cluster. This script is designed to manage distributed training resources.

```shell
# Example usage: sbatch ./scripts/slurm/pretrain_neurojepa.sh <your_config_name>

```

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

### Print Input Image Shape

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Prints the shape of the input image batch after preprocessing and loading. This helps verify the dimensions of the data being fed into the model.

```python
print(f"Shape for current scan: {x_img.shape}")
```

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

### MoE Block with Shared Experts in Neuro-JEPA

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Implements a Mixture-of-Experts (MoE) layer. It features a gating mechanism to route inputs to experts and includes shared experts for broader learning.

```python
      (mlp): MoE(
        (gate): Gate(
          (gate): Linear(in_features=768, out_features=16, bias=False)
        )
        (experts): ModuleList(
          (0-15): 16 x MLP(
            (fc1): Linear(in_features=768, out_features=384, bias=True)
            (act): GELU(approximate='none')
            (fc2): Linear(in_features=384, out_features=768, bias=True)
            (drop): Dropout(p=0.0, inplace=False)
          )
        )
        (shared_experts): MLP(
          (fc1): Linear(in_features=768, out_features=768, bias=True)
          (act): GELU(approximate='none')
          (fc2): Linear(in_features=768, out_features=768, bias=True)
          (drop): Dropout(p=0.0, inplace=False)
        )
      )
```

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

### MoE Layer with Gate in Neuro-JEPA

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Shows the implementation of a Mixture-of-Experts (MoE) layer, including its gating mechanism and expert MLPs. This layer dynamically selects experts based on input.

```python
MoE(
            (gate): Gate(
              (gate): Linear(in_features=768, out_features=16, bias=False)
            )
            (experts): ModuleList(
              (0-15): 16 x MLP(
                (fc1): Linear(in_features=768, out_features=384, bias=True)
                (act): GELU(approximate='none')
                (fc2): Linear(in_features=384, out_features=768, bias=True)
                (drop): Dropout(p=0.0, inplace=False)
              )
            )
            (shared_experts): MLP(
              (fc1): Linear(in_features=768, out_features=768, bias=True)
              (act): GELU(approximate='none')
              (fc2): Linear(in_features=768, out_features=768, bias=True)
              (drop): Dropout(p=0.0, inplace=False)
            )
          )
```

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

### Transformer Block with Mixture-of-Experts (MoE)

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

A transformer block incorporating a Mixture-of-Experts (MoE) layer for conditional computation. This allows for dynamic routing of inputs to different expert networks.

```python
    (5): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MoE(
        (gate): Gate(
          (gate): Linear(in_features=768, out_features=16, bias=False)
        )
        (experts): ModuleList(
          (0-15): 16 x MLP(
            (fc1): Linear(in_features=768, out_features=384, bias=True)
            (act): GELU(approximate='none')
            (fc2): Linear(in_features=384, out_features=768, bias=True)
            (drop): Dropout(p=0.0, inplace=False)
          )
        )
        (shared_experts): MLP(
          (fc1): Linear(in_features=768, out_features=768, bias=True)
          (act): GELU(approximate='none')
          (fc2): Linear(in_features=768, out_features=768, bias=True)
          (drop): Dropout(p=0.0, inplace=False)
        )
      )
    )
```

```python
    (7): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MoE(
        (gate): Gate(
          (gate): Linear(in_features=768, out_features=16, bias=False)
        )
        (experts): ModuleList(
          (0-15): 16 x MLP(
            (fc1): Linear(in_features=768, out_features=384, bias=True)
            (act): GELU(approximate='none')
            (fc2): Linear(in_features=384, out_features=768, bias=True)
            (drop): Dropout(p=0.0, inplace=False)
          )
        )
        (shared_experts): MLP(
          (fc1): Linear(in_features=768, out_features=768, bias=True)
          (act): GELU(approximate='none')
          (fc2): Linear(in_features=768, out_features=768, bias=True)
          (drop): Dropout(p=0.0, inplace=False)
        )
      )
    )
```

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

### VisionTransformer Block with MoE (continued)

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Details a transformer block that uses a Mixture-of-Experts (MoE) layer. This configuration is used in specific layers of the VisionTransformer.

```python
(11): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MoE(
        (gate): Gate(
          (gate): Linear(in_features=768, out_features=16, bias=False)
        )
        (experts): ModuleList(
          (0-15): 16 x MLP(
            (fc1): Linear(in_features=768, out_features=384, bias=True)
            (act): GELU(approximate='none')
            (fc2): Linear(in_features=384, out_features=768, bias=True)
            (drop): Dropout(p=0.0, inplace=False)
          )
        )
        (shared_experts): MLP(
          (fc1): Linear(in_features=768, out_features=768, bias=True)
          (act): GELU(approximate='none')
          (fc2): Linear(in_features=768, out_features=768, bias=True)
          (drop): Dropout(p=0.0, inplace=False)
        )
      )
    )
```

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

### Extract Features with Backbone

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Disables cuDNN for deterministic behavior and enables automatic mixed precision (AMP) for faster computation. It then passes the preprocessed image batch through the Neuro-JEPA backbone to extract features and MoE scores.

```python
torch.backends.cudnn.enabled = False
# Model output last layer and MoE scores per layer by default
with torch.autocast("cuda"):
    outs_features, outs_moe_scores = backbone(x_img)
```

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

### Standard MLP Block in Neuro-JEPA

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Represents a standard transformer block with LayerNorm, RoPEAttention, and an MLP. This is a common building block in the Neuro-JEPA architecture.

```python
Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLP(
        (fc1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU(approximate='none')
        (fc2): Linear(in_features=3072, out_features=768, bias=True)
        (drop): Dropout(p=0.0, inplace=False)
      )
    )
```

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

### MoE Block in Neuro-JEPA

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Illustrates a transformer block utilizing a Mixture-of-Experts (MoE) layer instead of a standard MLP. This block includes LayerNorm, RoPEAttention, and the MoE component.

```python
Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MoE(
        (gate): Gate(
          (gate): Linear(in_features=768, out_features=16, bias=False)
        )
        (experts): ModuleList(
          (0-15): 16 x MLP(
            (fc1): Linear(in_features=768, out_features=384, bias=True)
            (act): GELU(approximate='none')
            (fc2): Linear(in_features=384, out_features=768, bias=True)
            (drop): Dropout(p=0.0, inplace=False)
          )
        )
        (shared_experts): MLP(
          (fc1): Linear(in_features=768, out_features=768, bias=True)
          (act): GELU(approximate='none')
          (fc2): Linear(in_features=768, out_features=768, bias=True)
          (drop): Dropout(p=0.0, inplace=False)
        )
      )
    )
```

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

### VisionTransformer Block with MLP (continued)

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Details a transformer block featuring a standard MLP layer, following the attention mechanism. This is part of the VisionTransformer model.

```python
(10): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLP(
        (fc1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU(approximate='none')
        (fc2): Linear(in_features=3072, out_features=768, bias=True)
        (drop): Dropout(p=0.0, inplace=False)
      )
    )
```

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

### RoPEAttention Layer in Neuro-JEPA

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Details the RoPEAttention (Rotary Positional Embedding Attention) layer implementation. This includes query, key, value projections, attention dropout, output projection, and an attention gate.

```python
attn): RoPEAttention(
            (qkv): Linear(in_features=768, out_features=2304, bias=True)
            (attn_drop): Dropout(p=0.0, inplace=False)
            (proj): Linear(in_features=768, out_features=768, bias=True)
            (proj_drop): Dropout(p=0.0, inplace=False)
            (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
          )
```

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

### Print Output Feature Shape

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Prints the shape of the extracted features from the last layer of the backbone. This confirms the dimensions of the feature representation generated by the model.

```python
print(f"Shape of last layer output feature: {outs_features.shape}")
```

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

### Standard Transformer Block

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Represents a standard transformer block with Layer Normalization, RoPEAttention, and an MLP. Used in various stages of the model.

```python
    (4): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLP(
        (fc1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU(approximate='none')
        (fc2): Linear(in_features=3072, out_features=768, bias=True)
        (drop): Dropout(p=0.0, inplace=False)
      )
    )
```

```python
    (6): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLP(
        (fc1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU(approximate='none')
        (fc2): Linear(in_features=3072, out_features=768, bias=True)
        (drop): Dropout(p=0.0, inplace=False)
      )
    )
```

```python
    (8): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLP(
        (fc1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU(approximate='none')
        (fc2): Linear(in_features=3072, out_features=768, bias=True)
      )
```

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

### VisionTransformer Block with MoE

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Represents a transformer block utilizing a Mixture-of-Experts (MoE) layer instead of a standard MLP. This allows for conditional computation.

```python
(9): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MoE(
        (gate): Gate(
          (gate): Linear(in_features=768, out_features=16, bias=False)
        )
        (experts): ModuleList(
          (0-15): 16 x MLP(
            (fc1): Linear(in_features=768, out_features=384, bias=True)
            (act): GELU(approximate='none')
            (fc2): Linear(in_features=384, out_features=768, bias=True)
            (drop): Dropout(p=0.0, inplace=False)
          )
        )
        (shared_experts): MLP(
          (fc1): Linear(in_features=768, out_features=768, bias=True)
          (act): GELU(approximate='none')
          (fc2): Linear(in_features=768, out_features=768, bias=True)
          (drop): Dropout(p=0.0, inplace=False)
        )
      )
    )
```

```python
(1): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MoE(
        (gate): Gate(
          (gate): Linear(in_features=768, out_features=16, bias=False)
        )
        (experts): ModuleList(
          (0-15): 16 x MLP(
            (fc1): Linear(in_features=768, out_features=384, bias=True)
            (act): GELU(approximate='none')
            (fc2): Linear(in_features=384, out_features=768, bias=True)
            (drop): Dropout(p=0.0, inplace=False)
          )
        )
        (shared_experts): MLP(
          (fc1): Linear(in_features=768, out_features=768, bias=True)
          (act): GELU(approximate='none')
          (fc2): Linear(in_features=768, out_features=768, bias=True)
          (drop): Dropout(p=0.0, inplace=False)
        )
      )
    )
```

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

### RoPEAttention Block in Neuro-JEPA

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Details the RoPEAttention mechanism, which uses Rotary Positional Embeddings. It includes query, key, value projections, attention dropout, and output projection.

```python
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
```

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

### MLP Block in Neuro-JEPA

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Represents a standard Multi-Layer Perceptron block within the Neuro-JEPA architecture. It includes linear transformations, GELU activation, and dropout.

```python
    (mlp): MLP(
        (fc1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU(approximate='none')
        (fc2): Linear(in_features=3072, out_features=768, bias=True)
        (drop): Dropout(p=0.0, inplace=False)
      )
```

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

### SLURM Cooldown Script

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Shell script for managing a cooldown period on a SLURM cluster. This script might be used for resource cleanup or job termination.

```shell
# Example usage: sbatch ./scripts/slurm/cooldown_neurojepa.sh

```

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

### VisionTransformer Block with MLP

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Represents a standard transformer block containing Layer Normalization, RoPEAttention, and a standard MLP. Used in the VisionTransformer architecture.

```python
(0): Block(
      (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (attn): RoPEAttention(
        (qkv): Linear(in_features=768, out_features=2304, bias=True)
        (attn_drop): Dropout(p=0.0, inplace=False)
        (proj): Linear(in_features=768, out_features=768, bias=True)
        (proj_drop): Dropout(p=0.0, inplace=False)
        (proj_attn_gate): Linear(in_features=768, out_features=12, bias=True)
      )
      (drop_path): Identity()
      (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
      (mlp): MLP(
        (fc1): Linear(in_features=768, out_features=3072, bias=True)
        (act): GELU(approximate='none')
        (fc2): Linear(in_features=3072, out_features=768, bias=True)
        (drop): Dropout(p=0.0, inplace=False)
      )
    )
```

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

### Shared Experts MLP in Neuro-JEPA

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Defines the shared experts Multi-Layer Perceptron (MLP) within the Neuro-JEPA architecture. This component is used in conjunction with expert layers.

```python
shared_experts): MLP(
            (fc1): Linear(in_features=768, out_features=768, bias=True)
            (act): GELU(approximate='none')
            (fc2): Linear(in_features=768, out_features=768, bias=True)
            (drop): Dropout(p=0.0, inplace=False)
          )
```

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

### Neuro-JEPA Citation

Source: https://github.com/nyumedml/neuro-jepa/blob/main/README.md

Use this BibTeX entry to cite the Neuro-JEPA paper in your research. Ensure proper attribution when using the repository or its findings.

```bibtex
@misc{huang2026learningsparselatentpredictive,
      title={Learning Sparse Latent Predictive Foundation Model for Multimodal Neuroimaging},
      author={Haoxu Huang and Long Chen and Jingyun Chen and Jinu Hyun and James Ryan Loftus and Kara Melmed and Daniel Orringer and Jennifer Frontera and Seena Dehkharghani and Arjun Masurkar and Narges Razavian},
      year={2026},
      eprint={2606.14957},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2606.14957},
}
```

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

### VisionTransformer Patch Embedding

Source: https://github.com/nyumedml/neuro-jepa/blob/main/notebooks/extract_feature_example.ipynb

Defines the initial patch embedding layer using a 3D convolution. This layer processes input volumes into feature maps.

```python
  (patch_embed): PatchEmbed3D(
    (proj): Conv3d(1, 768, kernel_size=(12, 12, 12), stride=(12, 12, 12))
  )
```

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