### Install GeneJEPA using uv

Source: https://github.com/biostateai/genejepa/blob/main/README.md

Installs the GeneJEPA package and its dependencies using uv, a fast Python package installer. This is the recommended installation method for optimal performance.

```bash
uv sync
```

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

### Install GeneJEPA and Dependencies

Source: https://context7.com/biostateai/genejepa/llms.txt

Installs GeneJEPA using pip in an editable mode and includes optional dependencies for visualizations. It also shows how to use 'uv sync' for dependency management.

```bash
python -m venv .venv && source .venv/bin/activate
pip install -U pip
pip install -e .
# Optional visualizations used by callbacks
pip install umap-learn

# Using uv (recommended)
uv sync
```

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

### Initialize GeneJepa Model and Load Checkpoint

Source: https://context7.com/biostateai/genejepa/llms.txt

Demonstrates how to initialize the JepaLightningModule with various configurations, access the underlying model, simulate a training step, and load a pre-trained model from a checkpoint for inference. It shows how to get embeddings using the loaded model.

```python
from genejepa.train import JepaLightningModule
import torch

# Create configurations
model_config = ModelConfig(gene_vocab_size=60000)
train_config = TrainingConfig(
    learning_rate=1e-4,
    max_epochs=50,
    sim_coeff=1.0,
    var_coeff=25.0,
    cov_coeff=1.0,
)
exp_config = ExperimentConfig()
total_steps = 100000  # Calculated from data size

# Create Lightning module
module = JepaLightningModule(
    model_config=model_config,
    train_config=train_config,
    exp_config=exp_config,
    total_steps=total_steps,
)

# Access the underlying model
jepa_model = module.model

# Training step processes a batch and returns loss dict
batch = {
    "indices": torch.randint(0, 60000, (1000,)),
    "values": torch.randn(1000),
    "offsets": torch.tensor([0, 500, 1000]),
}
# output = module.training_step(batch, 0)

# Load from checkpoint
module = JepaLightningModule.load_from_checkpoint(
    "checkpoints/gene_jepa_tahoe/last.ckpt"
)
model = module.model.eval()

# Get embeddings for inference
with torch.no_grad():
    embeddings = model.get_embedding(
        batch["indices"],
        batch["values"],
        batch["offsets"],
        use_teacher=True
    )
```

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

### Install GeneJEPA using pip

Source: https://github.com/biostateai/genejepa/blob/main/README.md

Installs the GeneJEPA package and its optional visualization dependencies using pip. This method is suitable for editable installs and requires Python 3.11+ and PyTorch 2.6+.

```bash
python -m venv .venv && source .venv/bin/activate
pip install -U pip
pip install -e .
# Optional visualizations used by callbacks
pip install umap-learn
```

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

### Train GeneJEPA Model

Source: https://context7.com/biostateai/genejepa/llms.txt

Starts the training process for GeneJEPA on the Tahoe-100M dataset using PyTorch Lightning. Training automatically detects available devices and saves checkpoints. Weights & Biases logging is enabled if configured.

```bash
# Start training (Lightning auto-detects devices)
uv run -m genejepa.train

# Checkpoints are saved under checkpoints/gene_jepa_tahoe/
# W&B logging is enabled if wandb is configured
```

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

### GenePerceiverJEPA: Full JEPA Model with torch

Source: https://context7.com/biostateai/genejepa/llms.txt

Implements the complete JEPA model for self-supervised learning, including student and EMA teacher encoders, and an MLP predictor. It handles masked prediction tasks. The forward pass returns predicted, target, and student context representations. Includes methods for updating the teacher and getting inference embeddings.

```python
import torch
from genejepa.models import GenePerceiverJEPA
from genejepa.configs import ModelConfig

# Create full JEPA model
config = ModelConfig(
    gene_vocab_size=60000,
    d=768,
    latents_L=512,
    blocks_D=24,
    heads_h=12,
    mask_ratio=0.45,
    num_targets=1,
    min_context_genes=512,
    min_target_genes_per_block=16,
)
model = GenePerceiverJEPA(config)

# Training forward pass with automatic masking
indices = torch.randint(0, 60000, (5000,), dtype=torch.long)
values = torch.randn(5000, dtype=torch.float32)
offsets = torch.tensor([0, 1000, 2500, 3500, 5000], dtype=torch.long)

# Forward returns (predicted, target, student_context) representations
pred, target, student_ctx = model(indices, values, offsets)
print(f"Predicted: {pred.shape}, Target: {target.shape}, Context: {student_ctx.shape}")

# Update EMA teacher after gradient step
model.update_teacher()

# Inference: get stable embeddings using EMA teacher
model.eval()
with torch.no_grad():
    embeddings = model.get_embedding(indices, values, offsets, use_teacher=True)
print(f"Inference embeddings: {embeddings.shape}")
```

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

### Inference: Get Cell Embeddings

Source: https://github.com/biostateai/genejepa/blob/main/README.md

Loads a trained GeneJEPA model checkpoint and generates embeddings for cells. It takes ragged tensor inputs representing gene indices, expression values, and offsets, and outputs embeddings of shape [batch, d].

```python
import torch
from lightning.pytorch import seed_everything
from genejepa.train import JepaLightningModule

# Load your checkpoint (path under `checkpoints/gene_jepa_tahoe/`)
ckpt_path = "checkpoints/gene_jepa_tahoe/last.ckpt"
module = JepaLightningModule.load_from_checkpoint(ckpt_path)
model = module.model.eval()  # GenePerceiverJEPA

# Ragged inputs for a small batch of cells
# indices: concat of token_ids; values: log1p-standardized expression; offsets: prefix sums
indices = torch.tensor([10, 42, 7, 3, 9,  1, 2], dtype=torch.long)
values  = torch.tensor([0.1, 0.5, 0.3, 2.1, -0.2, 0.0, 1.1], dtype=torch.float32)
offsets = torch.tensor([0, 5, 7], dtype=torch.long)  # 2 samples: 5 tokens, then 2 tokens

with torch.no_grad():
    emb = model.get_embedding(indices, values, offsets, use_teacher=True)  # [batch, d] 
print(emb.shape)
```

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

### Create Training Configuration for Hyperparameters

Source: https://context7.com/biostateai/genejepa/llms.txt

Sets up training hyperparameters including loss coefficients for VICReg, optimizer settings like learning rate and weight decay, and training schedule parameters such as maximum epochs and gradient accumulation.

```python
from genejepa.configs import TrainingConfig

# Create training configuration
train_config = TrainingConfig(
    # VICReg loss coefficients
    sim_coeff=1.0,                   # Similarity (invariance) loss weight
    var_coeff=25.0,                  # Variance loss weight
    cov_coeff=1.0,                   # Covariance loss weight

    # Optimizer configuration
    learning_rate=1e-4,              # Peak learning rate
    weight_decay=2e-4,               # AdamW weight decay
    warmup_ratio=0.05,               # Warmup ratio of total steps
    adam_betas=(0.9, 0.98),          # Adam beta parameters

    # Training schedule
    max_epochs=50,                   # Maximum training epochs
    accumulate_grad_batches=2,       # Gradient accumulation steps
    gradient_clip_val=1.0,           # Gradient clipping value
)

print(f"Learning rate: {train_config.learning_rate}")
print(f"Max epochs: {train_config.max_epochs}")
print(f"Effective batch multiplier: {train_config.accumulate_grad_batches}")
```

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

### JepaLightningModule: PyTorch Lightning Training Module

Source: https://context7.com/biostateai/genejepa/llms.txt

A PyTorch Lightning module that wraps the JEPA model. It includes training logic, VICReg loss computation, EMA scheduling, and optimizer configuration with cosine learning rate decay. This module facilitates the training process within the PyTorch Lightning framework.

```python
import torch
import lightning as L
from genejepa.train import JepaLightningModule
from genejepa.configs import ModelConfig, TrainingConfig, ExperimentConfig

# Example instantiation (requires full config setup)
# model_config = ModelConfig(...)
# training_config = TrainingConfig(...)
# exp_config = ExperimentConfig(...)
# module = JepaLightningModule(model_config, training_config, exp_config)

```

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

### Create Experiment Configuration for Experiment Settings

Source: https://context7.com/biostateai/genejepa/llms.txt

Defines experiment-level settings such as the checkpoint directory, random seed, logging frequency, and Weights & Biases (W&B) integration parameters including project name, entity, and run name.

```python
from genejepa.configs import ExperimentConfig

# Create experiment configuration
exp_config = ExperimentConfig(
    checkpoint_dir="checkpoints/gene_jepa_tahoe",  # Checkpoint save directory
    random_seed=42,                                 # Global random seed
    log_every_n_steps=10,                           # Logging frequency
    wandb_project="genejepa",                       # W&B project name
    wandb_entity=None,                              # W&B entity (optional)
    wandb_run_name="GeneJEPA-Tahoe-100M",           # W&B run name
    validation_num_batches=4,                       # Batches for UMAP/metrics
    validation_plot_every_n_epochs=1,               # UMAP plot frequency
)

print(f"Checkpoint directory: {exp_config.checkpoint_dir}")
print(f"W&B project: {exp_config.wandb_project}")
```

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

### Create Data Configuration for Data Loading

Source: https://context7.com/biostateai/genejepa/llms.txt

Configures data loading parameters for the Tahoe-100M dataset. Includes settings for batch size per GPU, number of DataLoader workers, and the total number of training and validation samples per epoch.

```python
from genejepa.configs import DataConfig

# Create data configuration
data_config = DataConfig(
    batch_size=92,                   # Samples per GPU per step
    num_workers=8,                   # DataLoader workers
    train_samples=1_000_000,         # Training samples per epoch
    val_samples=10_000,              # Validation samples per epoch
)

print(f"Batch size: {data_config.batch_size}")
print(f"Training samples: {data_config.train_samples:,}")
```

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

### Complete GeneJepa Training Pipeline

Source: https://context7.com/biostateai/genejepa/llms.txt

Sets up a full end-to-end training pipeline for GeneJepa, including distributed training, logging with Wandb, and custom callbacks for validation. It configures data modules, model, logger, callbacks, and the PyTorch Lightning trainer.

```python
import os
import torch
import lightning as L
from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch.strategies import DDPStrategy
from dataclasses import asdict

from genejepa.configs import ModelConfig, TrainingConfig, DataConfig, ExperimentConfig
from genejepa.data import Tahoe100MDataModule
from genejepa.train import JepaLightningModule
from genejepa.callbacks import EmbeddingQualityValidator, SupervisedValidatorCallback

# Initialize configurations
model_config = ModelConfig()
train_config = TrainingConfig()
data_config = DataConfig()
exp_config = ExperimentConfig()

# Set global seed
L.seed_everything(exp_config.random_seed, workers=True)

# Setup data module
datamodule = Tahoe100MDataModule(data_config, exp_config)
datamodule.prepare_data()
datamodule.setup("fit")

# Finalize model config with vocabulary size from data
model_config.gene_vocab_size = datamodule.gene_vocab_size

# Calculate total training steps
num_devices = torch.cuda.device_count() if torch.cuda.is_available() else 1
steps_per_epoch = data_config.train_samples // (data_config.batch_size * num_devices)
total_steps = steps_per_epoch * train_config.max_epochs

# Create Lightning module
module = JepaLightningModule(model_config, train_config, exp_config, total_steps)

# Configure logger
logger = WandbLogger(
    project=exp_config.wandb_project,
    name=exp_config.wandb_run_name,
    config={
        "model": asdict(model_config),
        "training": asdict(train_config),
        "data": asdict(data_config),
    }
)

# Configure callbacks
callbacks = [
    ModelCheckpoint(
        dirpath=exp_config.checkpoint_dir,
        monitor="val_loss",
        save_top_k=2,
        mode="min",
        save_last=True,
    ),
    LearningRateMonitor(logging_interval="step"),
    EmbeddingQualityValidator(
        num_batches=exp_config.validation_num_batches,
        plot_every_n_epochs=exp_config.validation_plot_every_n_epochs,
    ),
    SupervisedValidatorCallback(
        foundation_gene_map=datamodule.gene_map,
        embedding_dim=model_config.d,
        global_mean=datamodule.global_mean,
        global_std=datamodule.global_std,
        probe_dataset_path="scvi-tools/human-lung-cell-atlas-scanvi",
        probe_cell_type_col="scanvi_label",
    ),
]

# Configure trainer
trainer = L.Trainer(
    accelerator="auto",
    devices="auto",
    strategy=DDPStrategy(find_unused_parameters=True) if num_devices > 1 else "auto",
    precision="bf16-mixed" if torch.cuda.is_bf16_supported() else "16-mixed",
    max_epochs=train_config.max_epochs,
    logger=logger,
    callbacks=callbacks,
    accumulate_grad_batches=train_config.accumulate_grad_batches,
    gradient_clip_val=train_config.gradient_clip_val,
)

# Start training
trainer.fit(module, datamodule=datamodule)
```

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

### Login to Hugging Face Hub

Source: https://github.com/biostateai/genejepa/blob/main/README.md

Authenticates with the Hugging Face Hub using the CLI. This is necessary for automatically pulling Tahoe-100M manifests required for training and other operations.

```bash
huggingface-cli login  # or set HUGGINGFACE_HUB_TOKEN
```

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

### Tahoe100MDataModule: Data Loading Pipeline with lightning

Source: https://context7.com/biostateai/genejepa/llms.txt

Handles streaming the Tahoe-100M dataset from Hugging Face Hub. It includes automatic downloading, gene mapping, and global normalization statistics computation. The module is designed to be used with PyTorch Lightning. It prepares data and sets up necessary configurations for training.

```python
import lightning as L
from genejepa.data import Tahoe100MDataModule
from genejepa.configs import DataConfig, ExperimentConfig

# Create data configuration
data_config = DataConfig(
    batch_size=92,
    num_workers=8,
    train_samples=1_000_000,
    val_samples=10_000,
)
exp_config = ExperimentConfig(random_seed=42)

# Initialize data module
datamodule = Tahoe100MDataModule(data_config, exp_config)

# Prepare data (downloads files on rank 0)
datamodule.prepare_data()

# Setup (loads gene map and computes global stats)
datamodule.setup("fit")

# Access properties after setup
print(f"Gene vocabulary size: {datamodule.gene_vocab_size}")
print(f"Global mean: {datamodule.global_mean:.4f}")
print(f"Global std: {datamodule.global_std:.4f}")

# Get data loaders
train_loader = datamodule.train_dataloader()
val_loader = datamodule.val_dataloader()

# Iterate over batches
for batch in train_loader:
    print(f"Batch keys: {batch.keys()}")
    print(f"Indices shape: {batch['indices'].shape}")
    print(f"Values shape: {batch['values'].shape}")
    print(f"Offsets shape: {batch['offsets'].shape}")
    break
```

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

### Train Probe on Frozen Embeddings with PyTorch

Source: https://context7.com/biostateai/genejepa/llms.txt

This snippet demonstrates how to train a linear probe on frozen embeddings using PyTorch. It initializes an Adam optimizer and a CrossEntropyLoss function, generates random labels, calculates the loss, performs backpropagation, and updates the optimizer.

```python
optimizer = torch.optim.Adam(probe.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

labels = torch.randint(0, num_classes, (32,))
loss = loss_fn(logits, labels)
loss.backward()
optimizer.step()
```

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

### Train GeneJEPA model

Source: https://github.com/biostateai/genejepa/blob/main/README.md

Initiates the training process for the GeneJEPA model on a single node. The script automatically detects available devices and saves checkpoints under `checkpoints/gene_jepa_tahoe/`. Weights & Biases logging is enabled if configured.

```bash
uv run -m genejepa.train
```

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

### Initialize scRNATokenizer with Model Configuration

Source: https://context7.com/biostateai/genejepa/llms.txt

Initializes the scRNATokenizer, which converts gene indices and expression values into embeddings. It uses a learnable projection to combine gene identity embeddings with Fourier-encoded expression values, requiring a ModelConfig object.

```python
import torch
from genejepa.tokenizer import scRNATokenizer
from genejepa.configs import ModelConfig

# Create tokenizer with model configuration
config = ModelConfig(gene_vocab_size=60000, d=768)
tokenizer = scRNATokenizer(config)
```

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

### Tokenize Gene Expression Data with torch

Source: https://context7.com/biostateai/genejepa/llms.txt

Demonstrates how to tokenize gene expression data using provided indices and values. The tokenizer combines identity, value, and projected embeddings. It takes gene token IDs and normalized expression values as input and outputs token embeddings.

```python
import torch

# indices: gene token IDs, values: normalized expression values
indices = torch.tensor([0, 1, 2, 100, 5000], dtype=torch.long)
values = torch.tensor([0.5, 1.2, -0.3, 2.1, 0.0], dtype=torch.float32)

# Forward pass produces token embeddings
# Assuming 'tokenizer' is an instance of a tokenizer class
# tokens = tokenizer(indices, values)
# print(f"Token embeddings shape: {tokens.shape}")  # [5, 768]
```

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

### Create Model Configuration with Custom Parameters

Source: https://context7.com/biostateai/genejepa/llms.txt

Defines the architecture and masking parameters for the GeneJEPA model. It includes settings for embedding dimensions, Perceiver layers, transformer blocks, attention heads, masking ratios, and EMA teacher configurations.

```python
from genejepa.configs import ModelConfig

# Create configuration with custom parameters
config = ModelConfig(
    d=768,                           # Embedding dimension
    latents_L=512,                   # Number of latent tokens in Perceiver
    blocks_D=24,                     # Number of transformer blocks
    heads_h=12,                      # Number of attention heads
    cross_attn_chunk_size=32,        # Chunk size for memory-efficient attention
    gene_vocab_size=60_000,          # Gene vocabulary size (set from datamodule)

    # Masking configuration
    mask_ratio=0.45,                 # Ratio of genes to mask as targets
    num_targets=1,                   # Number of target blocks
    min_context_genes=512,           # Minimum genes in context
    min_target_genes_per_block=16,   # Minimum genes per target block

    # EMA teacher configuration
    ema_start_decay=0.992,           # Initial EMA decay rate
    ema_end_decay=0.9995,            # Final EMA decay rate
    ema_warmup_steps=2000,           # Steps before EMA updates begin

    # Tokenizer configuration
    identity_value_split_ratio=0.5,  # Ratio of dimensions for identity vs value
    fourier_num_frequencies=64,      # Number of Fourier frequencies
    fourier_min_freq=0.1,            # Minimum frequency
    fourier_max_freq=100.0,          # Maximum frequency
    fourier_freq_scale=1.0,          # Frequency scaling factor

    # Predictor configuration
    predictor_depth=3,               # Number of predictor layers
    predictor_expansion_factor=4,    # MLP expansion factor
)

print(f"Model dimension: {config.d}")
print(f"Transformer blocks: {config.blocks_D}")
print(f"Attention heads: {config.heads_h}")
```

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

### ModelConfig: GeneJEPA Architecture Configuration

Source: https://context7.com/biostateai/genejepa/llms.txt

Defines the configuration for the GeneJEPA model architecture. This includes parameters such as embedding dimension, number of latent tokens, transformer blocks, and masking strategies for JEPA training.

```python
from genejepa.configs import ModelConfig

# Example usage (assuming ModelConfig is a class or has default values):
# config = ModelConfig(
#     embedding_dim=768,
#     num_latent_tokens=256,
#     num_transformer_blocks=12,
#     masking_ratio=0.75
# )
# print(config)
```

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

### Export Gene Vocabulary Map and Global Stats

Source: https://github.com/biostateai/genejepa/blob/main/README.md

Exports the foundation gene map and global statistics without initiating model training. This command requires specifying output paths for the map and stats, as well as the foundation metadata.

```bash
uv run -m genejepa.train \
  --export-foundation-map hf_data_cache/foundation_gene_map.parquet \
  --export-global-stats hf_data_cache/global_stats.json \
  --foundation-meta hf_data_cache/data/gene_metadata.parquet \
  --export-only
```

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

### SupervisedValidatorCallback for Linear Probing

Source: https://context7.com/biostateai/genejepa/llms.txt

This callback evaluates embedding quality by training a linear probe for cell type classification on external datasets. It processes datasets, extracts embeddings, trains an MLP classifier, and reports validation accuracy to W&B. It requires gene mapping, embedding dimensions, and dataset paths.

```python
from genejepa.callbacks import SupervisedValidatorCallback

validator = SupervisedValidatorCallback(
    foundation_gene_map=gene_map,
    embedding_dim=768,
    global_mean=0.8255,
    global_std=0.3135,
    probe_dataset_path="scvi-tools/human-lung-cell-atlas-scanvi",
    probe_cell_type_col="scanvi_label",
    probe_train_epochs=5,
    probe_batch_size=32,
    probe_lr=1e-3,
    run_every_n_epochs=1,
    max_probe_cells=50_000,
)
```

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

### EmbeddingQualityValidator Callback for Monitoring

Source: https://context7.com/biostateai/genejepa/llms.txt

The EmbeddingQualityValidator callback monitors embedding quality during training by calculating collapse metrics and generating UMAP visualizations. It logs these metrics and plots to Weights & Biases (W&B). Key metrics include average cosine similarity and embedding norm variance.

```python
from genejepa.callbacks import EmbeddingQualityValidator

validator = EmbeddingQualityValidator(
    num_batches=4,
    plot_every_n_epochs=1,
    color_by="drug",
)
```

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

### LinearProbeMLP for Downstream Task Evaluation

Source: https://context7.com/biostateai/genejepa/llms.txt

Implements a simple Multi-Layer Perceptron (MLP) classifier designed for evaluating the quality of learned embeddings on downstream tasks, specifically cell type classification. It takes embeddings as input and outputs class logits.

```python
import torch
import torch.nn as nn
from genejepa.callbacks import LinearProbeMLP

# Create probe for cell type classification
embedding_dim = 768
num_classes = 50  # Number of cell types

probe = LinearProbeMLP(embedding_dim, num_classes)

# Example: classify cell embeddings
embeddings = torch.randn(32, embedding_dim)  # Batch of cell embeddings
logits = probe(embeddings)
predictions = logits.argmax(dim=1)

print(f"Logits shape: {logits.shape}")  # [32, 50]
print(f"Predictions shape: {predictions.shape}")  # [32]
```

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

### GenePerceiverEncoder: Perceiver-Style Cell Encoder with torch

Source: https://context7.com/biostateai/genejepa/llms.txt

Encodes variable-length sets of gene tokens into fixed-size cell representations. It uses cross-attention from learned latent queries to gene tokens, followed by self-attention transformer blocks. Input is ragged batch data (indices, values, offsets), and output is fixed-size cell embeddings.

```python
import torch
from genejepa.models import GenePerceiverEncoder
from genejepa.configs import ModelConfig

# Create encoder with configuration
config = ModelConfig(
    gene_vocab_size=60000,
    d=768,
    latents_L=512,      # Number of latent queries
    blocks_D=24,        # Transformer depth
    heads_h=12,         # Attention heads
)
encoder = GenePerceiverEncoder(config)

# Prepare ragged batch input
# Cell 1: 3 genes, Cell 2: 4 genes
indices = torch.tensor([10, 20, 30, 100, 200, 300, 400], dtype=torch.long)
values = torch.tensor([0.5, 1.0, -0.2, 0.8, 1.5, 0.3, 2.0], dtype=torch.float32)
offsets = torch.tensor([0, 3, 7], dtype=torch.long)  # Defines cell boundaries

# Encode cells to fixed-size embeddings
with torch.no_grad():
    cell_embeddings = encoder(indices, values, offsets)
print(f"Cell embeddings shape: {cell_embeddings.shape}")  # [2, 768]
```

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