### Basic TorchEBM Example

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/index.md

Demonstrates creating a 2D Gaussian energy model, initializing a Langevin dynamics sampler, and generating samples. Use this to get started with TorchEBM.

```python
import torch
from torchebm.core import GaussianModel
from torchebm.samplers.langevin_dynamics import LangevinDynamics

# Create an energy model (2D Gaussian)
energy_fn = GaussianModel(
    mean=torch.zeros(2),
    cov=torch.eye(2)
)

# Create a sampler
sampler = LangevinDynamics(
    energy_function=energy_fn,
    step_size=0.01
)

# Generate samples
samples = sampler.sample(
    dim=2, n_steps=100, n_samples=1000
)

# Print sample statistics
print(f"Sample mean: {samples.mean(0)}")
print(f"Sample std: {samples.std(0)}")
```

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

### Run TorchEBM Examples

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/examples/index.md

Demonstrates how to clone the repository, set up the environment, list available examples, and run a specific example using the main.py script.

```bash
git clone https://github.com/soran-ghaderi/torchebm.git
cd torchebm
pip install -e .
python examples/main.py --list
python examples/main.py samplers/langevin/visualization_trajectory
```

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

### List All Available Examples

Source: https://github.com/soran-ghaderi/torchebm/blob/master/examples/README.md

Lists all executable example scripts available within the TorchEBM repository. This command helps in discovering available functionalities.

```bash
python examples/main.py --list
```

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

### Complete Training Example with Contrastive Divergence

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/loss_functions.md

A full example demonstrating the training of an Energy-Based Model using the Contrastive Divergence loss function. This includes model definition, dataset loading, sampler initialization, loss function setup, and the training loop.

```python
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader

from torchebm.core import BaseModel
from torchebm.samplers import LangevinDynamics
from torchebm.losses import ContrastiveDivergence
from torchebm.datasets import TwoMoonsDataset

class MLPModel(BaseModel):
    def __init__(self, input_dim, hidden_dim=64):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, 1),
            nn.Tanh(),
        )

    def forward(self, x):
        return self.network(x).squeeze(-1)

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

INPUT_DIM = 2
HIDDEN_DIM = 16
BATCH_SIZE = 256
EPOCHS = 100
LEARNING_RATE = 1e-3
CD_K = 10
USE_PCD = True

dataset = TwoMoonsDataset(n_samples=3000, noise=0.05, seed=42, device=device)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)

model = MLPModel(INPUT_DIM, HIDDEN_DIM).to(device)
sampler = LangevinDynamics(
    model=model,
    step_size=0.1,
    device=device,
)
loss_fn = ContrastiveDivergence(
    model=model,
    sampler=sampler,
    k_steps=CD_K,
    persistent=USE_PCD,
    buffer_size=BATCH_SIZE
).to(device)

optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

losses = []
print("Starting training...")
for epoch in range(EPOCHS):
    model.train()
    epoch_loss = 0.0

    for i, data_batch in enumerate(dataloader):
        optimizer.zero_grad()

        loss, negative_samples = loss_fn(data_batch)

        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()

        epoch_loss += loss.item()

    avg_epoch_loss = epoch_loss / len(dataloader)
    losses.append(avg_epoch_loss)
    print(f"Epoch [{epoch+1}/{EPOCHS}], Average Loss: {avg_epoch_loss:.4f}")

plt.figure(figsize=(10, 6))
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('docs/assets/images/loss_functions/cd_training_loss.png')
plt.show()
```

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

### Run Specific Example with TorchEBM

Source: https://github.com/soran-ghaderi/torchebm/blob/master/examples/README.md

Demonstrates how to execute a specific example script from the TorchEBM repository using the main entry point. Ensure you are in the repository root directory.

```bash
python examples/main.py samplers/langevin/gaussian_sampling
```

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

### Run Python Example Directly

Source: https://github.com/soran-ghaderi/torchebm/blob/master/examples/README.md

Shows how to execute a TorchEBM example script directly using the Python interpreter. This is an alternative to using the main entry point.

```bash
python examples/energy_functions/energy_functions/landscape_2d.py
```

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

### Install TorchEBM

Source: https://context7.com/soran-ghaderi/torchebm/llms.txt

Install the TorchEBM library using pip. Optionally, install torchdiffeq for adaptive ODE solvers.

```bash
pip install torchebm
# optional: adaptive ODE solvers
pip install torchdiffeq
```

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

### Clone and Set Up Development Environment

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/developer_guide/getting_started.md

Clone the repository, set up a Python virtual environment, and install development dependencies. Use `.[dev]` for general development and `.[docs]` for documentation work.

```bash
git clone https://github.com/<your-user>/torchebm.git
cd torchebm
python -m venv .venv
source .venv/bin/activate        # Windows: .venv\Scripts\activate
pip install -e ".[dev]"
```

```bash
pip install -e ".[docs]"
mkdocs serve                     # live preview at http://127.0.0.1:8000
```

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

### Train Generative Model with Equilibrium Matching

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/index.md

This example demonstrates training a generative model using Equilibrium Matching with a custom VelocityNet and sampling with FlowSampler. Ensure you have PyTorch and TorchEBM installed.

```python
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchebm.losses import EquilibriumMatchingLoss
from torchebm.samplers import FlowSampler
from torchebm.core import BaseModel
from torchebm.datasets import EightGaussiansDataset

dataset = EightGaussiansDataset(n_samples=8192)
loader = DataLoader(dataset, batch_size=256, shuffle=True)

# Any nn.Module with forward(x, t) works
class VelocityNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 256), nn.SiLU(),
                                 nn.Linear(256, 256), nn.SiLU(), nn.Linear(256, 2))
    def forward(self, x, t, **kwargs):
        return self.net(x)

model = VelocityNet()

loss_fn = EquilibriumMatchingLoss(
    model=model, interpolant="linear", energy_type="dot",
)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)

for epoch in range(50):
    for x in loader:
        loss = loss_fn(x)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# Sample via ODE flow
flow = FlowSampler(model=model, interpolant="linear", negate_velocity=True)
flow_samples = flow.sample_ode(torch.randn(1000, 2), num_steps=100)

```

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

### Using Torchebm Synthetic Datasets with DataLoader

Source: https://context7.com/soran-ghaderi/torchebm/llms.txt

Shows how to instantiate various 2D synthetic datasets and use them with PyTorch's DataLoader. Includes examples of getting all data as a tensor and seeded regeneration for reproducibility.

```python
import torch
from torch.utils.data import DataLoader
from torchebm.datasets import (
    GaussianMixtureDataset, EightGaussiansDataset, TwoMoonsDataset,
    SwissRollDataset, CircleDataset, CheckerboardDataset, PinwheelDataset,
)

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

datasets = {
    "GaussianMixture":  GaussianMixtureDataset(n_samples=2000, n_components=8, std=0.05, seed=42),
    "EightGaussians":   EightGaussiansDataset(n_samples=2000, std=0.02, seed=42),
    "TwoMoons":         TwoMoonsDataset(n_samples=2000, noise=0.05, seed=42),
    "SwissRoll":        SwissRollDataset(n_samples=2000, noise=0.05, seed=42),
    "Circle":           CircleDataset(n_samples=2000, noise=0.05, radius=1.0, seed=42),
    "Checkerboard":     CheckerboardDataset(n_samples=2000, range_limit=4.0, seed=42),
    "Pinwheel":         PinwheelDataset(n_samples=2000, n_classes=5, seed=42),
}

for name, ds in datasets.items():
    tensor = ds.get_data()              # entire dataset as a tensor
    print(f"{name:16s}: {tensor.shape}")

# Direct DataLoader integration
loader = DataLoader(datasets["TwoMoons"], batch_size=256, shuffle=True)
for batch in loader:
    print(f"Batch shape: {batch.shape}")   # torch.Size([256, 2])
    break

# Seeded regeneration for reproducibility
datasets["TwoMoons"].regenerate(seed=123)
```

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

### Install TorchEBM with CUDA Support

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/faq.md

Install TorchEBM using pip. Ensure PyTorch is installed with CUDA support beforehand.

```bash
pip install torchebm
```

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

### Implement a Custom Sampler in Python

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/samplers.md

Subclass BaseSampler and implement the sample() method to create a custom sampler. This example shows a basic structure with a custom_step method and parameter scheduling.

```python
from torchebm.core import BaseSampler, BaseModel
import torch
from typing import Optional, Union, Tuple, List, Dict

class MyCustomSampler(BaseSampler):
    def __init__(
        self,
        model: BaseModel,
        my_parameter: float = 0.1,
        dtype: torch.dtype = torch.float32,
        device: Optional[Union[str, torch.device]] = None,
    ):
        super().__init__(model=model, dtype=dtype, device=device)
        self.my_parameter = my_parameter
        
        self.register_scheduler("my_parameter", ConstantScheduler(my_parameter))

    def custom_step(self, x: torch.Tensor) -> torch.Tensor:
        param_value = self.get_scheduled_value("my_parameter")

        gradient = self.model.gradient(x)

        noise = torch.randn_like(x)
        new_x = x - param_value * gradient + noise * 0.01

        return new_x
    
    @torch.no_grad()
    def sample(
        self,
        x: Optional[torch.Tensor] = None,
        dim: int = 10,
        n_steps: int = 100,
        n_samples: int = 1,
        thin: int = 1,
        return_trajectory: bool = False,
        return_diagnostics: bool = False,
        *args,
        **kwargs,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]:
        self.reset_schedulers()

        if x is None:
            x = torch.randn(n_samples, dim, dtype=self.dtype, device=self.device)
        else:
            x = x.to(self.device)

        if return_trajectory:
            trajectory = torch.empty(
                (n_samples, n_steps, dim), dtype=self.dtype, device=self.device
            )

        if return_diagnostics:
            diagnostics = self._setup_diagnostics(dim, n_steps, n_samples=n_samples)

        for i in range(n_steps):
            self.step_schedulers()

            x = self.custom_step(x)

            if return_trajectory:
                trajectory[:, i, :] = x

            if return_diagnostics:
                pass

        if return_trajectory:
            if return_diagnostics:
                return trajectory, diagnostics
            return trajectory
        if return_diagnostics:
            return x, diagnostics
        return x
    
    def _setup_diagnostics(self, dim: int, n_steps: int, n_samples: int = None) -> torch.Tensor:
        """Optional method to setup diagnostic storage"""
        if n_samples is not None:
            return torch.empty(
                (n_steps, 3, n_samples, dim), device=self.device, dtype=self.dtype
            )
        else:
            return torch.empty((n_steps, 3, dim), device=self.device, dtype=self.dtype)
```

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

### Numerical Integrator Examples

Source: https://context7.com/soran-ghaderi/torchebm/llms.txt

Illustrates the use of various numerical integrators including Euler-Maruyama, Heun, Leapfrog, and RK4 for SDEs, ODEs, and Hamiltonian dynamics. These can be used directly for custom dynamics.

```python
import torch
from torchebm.integrators import (
    EulerMaruyamaIntegrator, HeunIntegrator,
    LeapfrogIntegrator, RK4Integrator,
)

device = torch.device("cpu")
dim, batch = 2, 100
x = torch.randn(batch, dim)

# --- Euler-Maruyama for overdamped Langevin SDE ---
em = EulerMaruyamaIntegrator(device=device)
drift = lambda x_, t_: -x_     # OU process: drift toward origin
result = em.integrate(
    state={"x": x},
    step_size=0.01,
    n_steps=200,
    drift=drift,
    noise_scale=1.0,
)
print("EM final:", result["x"].shape)   # torch.Size([100, 2])

# --- Heun (2nd-order) for deterministic ODE ---
heun = HeunIntegrator(device=device)
result_h = heun.integrate(state={"x": x}, step_size=0.01, n_steps=200, drift=drift)
print("Heun final:", result_h["x"].shape)

# --- Leapfrog for Hamiltonian dynamics ---
lf = LeapfrogIntegrator(device=device)
p  = torch.randn_like(x)
result_lf = lf.integrate(
    state={"x": x, "p": p},
    step_size=0.1, n_steps=10,
    drift=lambda x_, t_: -x_,   # gradient of potential
    mass=None,
)
print("Leapfrog x:", result_lf["x"].shape, " p:", result_lf["p"].shape)

# --- RK4 for stiff ODEs ---
rk4 = RK4Integrator(device=device)
result_rk4 = rk4.integrate(state={"x": x}, step_size=0.05, n_steps=50, drift=drift)
print("RK4 final:", result_rk4["x"].shape)
```

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

### Basic Langevin Dynamics Sampling

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/blog/posts/langevin_dynamics_sampling.md

Use this snippet to sample from a 2D Gaussian distribution. Ensure TorchEBM and Matplotlib are installed. Samples are plotted for visualization.

```python
import torch
import matplotlib.pyplot as plt
from torchebm.core import GaussianEnergy
from torchebm.samplers.langevin_dynamics import LangevinDynamics

# Create energy function for a 2D Gaussian
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dim = 2  # dimension of the state space
n_steps = 100  # steps between samples
n_samples = 1000  # num of samples
mean = torch.tensor([1.0, -1.0])
cov = torch.tensor([[1.0, 0.5], [0.5, 2.0]])
energy_fn = GaussianEnergy(mean, cov, device=device)

# Initialize sampler
sampler = LangevinDynamics(
    energy_function=energy_fn,
    step_size=0.01,
    noise_scale=0.1,
    device=device,
)

# Generate samples
initial_state = torch.zeros(n_samples, dim, device=device)
samples = sampler.sample(
    x=initial_state,
    n_steps=n_steps,
    n_samples=n_samples,
)

# Plot results
samples = samples.cpu().numpy()
plt.figure(figsize=(10, 5))
plt.scatter(samples[:, 0], samples[:, 1], alpha=0.1)
plt.title("Samples from 2D Gaussian using Langevin Dynamics")
plt.xlabel("x₁")
plt.ylabel("x₂")
plt.show()
```

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

### Basic Contrastive Divergence Training

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/loss_functions.md

Demonstrates the fundamental setup for training an EBM using Contrastive Divergence with Langevin Dynamics for MCMC sampling. Requires defining a model, sampler, and loss function.

```python
import torch
import torch.nn as nn
import torch.optim as optim
from torchebm.core import BaseModel
from torchebm.losses import ContrastiveDivergence
from torchebm.samplers import LangevinDynamics

class MLPModel(BaseModel):
    def __init__(self, input_dim, hidden_dim=64):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, 1),
            nn.Tanh(),
        )

    def forward(self, x):
        return self.network(x).squeeze(-1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MLPModel(input_dim=2, hidden_dim=64).to(device)

sampler = LangevinDynamics(
    model=model,
    step_size=0.1,
    device=device
)

loss_fn = ContrastiveDivergence(
    model=model,
    sampler=sampler,
    k_steps=10
)

optimizer = optim.Adam(model.parameters(), lr=0.001)

data_batch = torch.randn(128, 2).to(device)
optimizer.zero_grad()
loss, negative_samples = loss_fn(data_batch)
loss.backward()
optimizer.step()
```

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

### Train EBM with Contrastive Divergence

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/training.md

Complete example demonstrating training an EBM using the Contrastive Divergence (CD) method with TorchEBM. Includes model definition, sampler, loss function, optimizer, and visualization.

```python
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
import os

from torchebm.core import BaseModel, CosineScheduler
from torchebm.samplers import LangevinDynamics
from torchebm.losses import ContrastiveDivergence
from torchebm.datasets import TwoMoonsDataset

torch.manual_seed(42)
np.random.seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed(42)

os.makedirs("training_plots", exist_ok=True)

INPUT_DIM = 2
HIDDEN_DIM = 16
BATCH_SIZE = 256
EPOCHS = 200
LEARNING_RATE = 1e-3

SAMPLER_STEP_SIZE = CosineScheduler(start_value=3e-2, end_value=5e-3, n_steps=100)
SAMPLER_NOISE_SCALE = CosineScheduler(start_value=3e-1, end_value=1e-2, n_steps=100)

CD_K = 10
USE_PCD = True
VISUALIZE_EVERY = 20

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

dataset = TwoMoonsDataset(n_samples=3000, noise=0.05, seed=42, device=device)
real_data_for_plotting = dataset.get_data()
dataloader = DataLoader(
    dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    drop_last=True,
)

model = MLPModel(INPUT_DIM, HIDDEN_DIM).to(device)
sampler = LangevinDynamics(
    model=model,
    step_size=SAMPLER_STEP_SIZE,
    noise_scale=SAMPLER_NOISE_SCALE,
    device=device,
)
loss_fn = ContrastiveDivergence(
    model=model,
    sampler=sampler,
    k_steps=CD_K,
    persistent=USE_PCD,
    buffer_size=BATCH_SIZE,
).to(device)

optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

losses = []
print("Starting training...")
for epoch in range(EPOCHS):
    model.train()
    epoch_loss = 0.0

    for i, data_batch in enumerate(dataloader):
        optimizer.zero_grad()

        loss, negative_samples = loss_fn(data_batch)

        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()

        epoch_loss += loss.item()

    avg_epoch_loss = epoch_loss / len(dataloader)
    losses.append(avg_epoch_loss)
    print(f"Epoch [{epoch+1}/{EPOCHS}], Average Loss: {avg_epoch_loss:.4f}")

    if (epoch + 1) % VISUALIZE_EVERY == 0 or epoch == 0:
        print("Generating visualization...")
        plot_energy_and_samples(
            model=model,
            real_samples=real_data_for_plotting,
            sampler=sampler,
            epoch=epoch + 1,
            device=device,
            plot_range=2.5,
            k_sampling=200,
        )

# Plot the training loss
plt.figure(figsize=(10, 6))
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss')
plt.grid(True, alpha=0.3)
plt.savefig('docs/assets/images/training/cd_training_loss.png')
plt.show()
```

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

### Set Up Device and Dataset

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/examples/training/index.md

Configure the computation device (CPU or GPU) and initialize the GaussianMixtureDataset for training. Create a DataLoader for batching.

```python
# Set up device and dataset
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dataset = GaussianMixtureDataset(
    n_samples=2048,
    n_components=8,
    std=0.1,
    radius=1.5,
    device=device,
    seed=42,
)
dataloader = DataLoader(dataset, batch_size=256, shuffle=True)
```

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

### Direct Sampler Usage with Langevin Dynamics

Source: https://context7.com/soran-ghaderi/torchebm/llms.txt

Demonstrates how to directly use a sampler with predefined step and noise schedulers. Schedulers advance automatically within the sampler's sample method.

```python
import torch
from torchebm.models import DoubleWellModel
from torchebm.schedulers import LinearScheduler, CosineScheduler
from torchebm.samplers import LangevinDynamics

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

model   = DoubleWellModel(device=device)
step_sched  = LinearScheduler(start_value=0.05, end_value=0.001, n_steps=500)
noise_sched = CosineScheduler(start_value=1.0, end_value=0.1, n_steps=500)

sampler = LangevinDynamics(
    model=model,
    step_size=step_sched,
    noise_scale=noise_sched,
    device=device,
)
# Schedulers advance automatically each MCMC step inside sampler.sample()
samples = sampler.sample(dim=2, n_samples=200, n_steps=500)
print(samples.shape)    # torch.Size([200, 2])
```

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

### Profile with All Options

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/developer_guide/profiling.md

A shortcut to enable tracing, top N operations, memory snapshots, and NVTX ranges for comprehensive profiling.

```bash
python benchmarks/profiler.py run --all --component <component_name>
```

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

### Set up Training Components for EBM

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/getting_started.md

Initializes the sampler, loss function, and optimizer for training an EBM. Requires importing necessary classes from torchebm and torch.optim.

```python
from torchebm.losses import ContrastiveDivergence
from torchebm.samplers import LangevinDynamics
from torch.optim import Adam

sampler = LangevinDynamics(
    energy_function=neural_energy_fn,
    step_size=10.0,
    noise_scale=0.1,
    n_steps=60
)

cd_loss = ContrastiveDivergence(sampler=sampler)

optimizer = Adam(neural_energy_fn.parameters(), lr=1e-4)
```

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

### Conventional Commit Examples

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/developer_guide/getting_started.md

Examples of commit messages following the Conventional Commits specification. Use types like `feat`, `fix`, `perf`, `refactor`, `test`, `docs`, `style`, `build`, `ci`, `chore`. Breaking changes require a `!` and a `BREAKING CHANGE:` footer.

```text
feat(samplers): add adaptive step size to LangevinDynamics
fix(losses): correct gradient sign in EquilibriumMatching
perf(integrators): cache RK buffers on device once per integrate()
docs(developer_guide): tighten profiling page
```

```text
refactor(core)!: rename BaseSampler.step -> BaseSampler.step_one

BREAKING CHANGE: external subclasses must rename the overridden method.
```

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

### Line-Level Profiling

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/developer_guide/profiling.md

Perform line-level profiling on a specific function within a module. Requires the `line_profiler` package to be installed.

```bash
python benchmarks/profiler.py run --line module:fn --component <component_name>
```

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

### Define a Custom EBM Model

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/training.md

Subclass `BaseModel` to create custom energy-based models. This example defines a simple MLP model.

```python
import torch
import torch.nn as nn
from torchebm.core import BaseModel

class MLPModel(BaseModel):
    def __init__(self, input_dim: int, hidden_dim: int = 64):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, 1),
            nn.Tanh(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x).squeeze(-1)
```

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

### Simple Langevin Dynamics Sampler

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/samplers.md

A basic implementation of Langevin dynamics for sampling. It requires a BaseModel and handles device placement and scheduling.

```python
class SimpleLangevin(BaseSampler):
    def __init__(
        self,
        model: BaseModel,
        step_size: float = 0.01,
        noise_scale: float = 1.0,
        dtype: torch.dtype = torch.float32,
        device: Optional[Union[str, torch.device]] = None,
    ):
        super().__init__(model=model, dtype=dtype, device=device)
        self.register_scheduler("step_size", ConstantScheduler(step_size))
        self.register_scheduler("noise_scale", ConstantScheduler(noise_scale))
    
    def langevin_step(self, x: torch.Tensor) -> torch.Tensor:
        step_size = self.get_scheduled_value("step_size")
        noise_scale = self.get_scheduled_value("noise_scale")
        
        gradient = self.model.gradient(x)
        noise = torch.randn_like(x)
        
        new_x = (
            x 
            - step_size * gradient 
            + torch.sqrt(torch.tensor(2.0 * step_size)) * noise_scale * noise
        )
        return new_x
    
    @torch.no_grad()
    def sample(
        self,
        x: Optional[torch.Tensor] = None,
        dim: int = 10,
        n_steps: int = 100,
        n_samples: int = 1,
        thin: int = 1,
        return_trajectory: bool = False,
        return_diagnostics: bool = False,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]:
        self.reset_schedulers()

        if x is None:
            x = torch.randn(n_samples, dim, dtype=self.dtype, device=self.device)
        else:
            x = x.to(self.device)

        if return_trajectory:
            trajectory = torch.empty(
                (n_samples, n_steps, dim), dtype=self.dtype, device=self.device
            )

        for i in range(n_steps):
            self.step_schedulers()
            x = self.langevin_step(x)

            if return_trajectory:
                trajectory[:, i, :] = x

        if return_trajectory:
            return trajectory
        return x
```

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

### Benchmark Execution Commands

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/developer_guide/benchmarking.md

Run the full benchmark suite, a quick smoke test, a specific module, or filter benchmarks by name. Options are available to include compile and AMP variants.

```bash
bash benchmarks/run.sh

# Smoke test (small scale only)
bash benchmarks/run.sh --quick

# One module
bash benchmarks/run.sh --module losses

# One benchmark by name (pytest -k)
bash benchmarks/run.sh --filter "ScoreMatching"

# With compile / AMP variants
bash benchmarks/run.sh --compile --amp
```

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

### Python Function with Docstring

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/developer_guide/code_guidelines.md

Example of a Python function adhering to Google-style docstrings with raw strings and inline LaTeX formatting for mathematical expressions.

```python
def sample_chain(
    dim: int,
    n_steps: int,
    n_samples: int = 1,
) -> tuple[torch.Tensor, dict]:
    r"""Run a Markov chain.

    Args:
        dim: Sample dimensionality.
        n_steps: Number of MCMC steps.
        n_samples: Number of parallel chains.

    Returns:
        Final samples and a diagnostics dict.
    """

```

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

### MCMC Sampling with Langevin Dynamics

Source: https://github.com/soran-ghaderi/torchebm/blob/master/README.md

Demonstrates how to perform MCMC sampling using Langevin Dynamics. Requires a model and sampler initialization.

```python
import torch
from torchebm.core import GaussianModel
from torchebm.samplers import LangevinDynamics

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GaussianModel(mean=torch.zeros(2), cov=torch.eye(2), device=device)

sampler = LangevinDynamics(model=model, step_size=0.01, device=device)
samples = sampler.sample(x=torch.randn(500, 2, device=device), n_steps=100)
print(samples.shape)  # torch.Size([500, 2])
```

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

### GPU Accelerated Sampling with DoubleWellModel

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/parallel_sampling.md

Leverage GPU acceleration for maximum performance in parallel sampling. This example uses the DoubleWellModel and measures execution time.

```python
import time
import torch
from torchebm.core import DoubleWellModel
from torchebm.samplers import LangevinDynamics

model = DoubleWellModel()

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

sampler = LangevinDynamics(
    model=model,
    step_size=0.01,
    device=device
)

n_samples = 50000
dim = 2

initial_points = torch.randn(n_samples, dim, device=device)

start_time = time.time()
samples = sampler.sample(
    x=initial_points,
    n_steps=1000,
    return_trajectory=False
)
end_time = time.time()

print(f"Generated {n_samples} samples in {end_time - start_time:.2f} seconds")
print(f"Average time per sample: {(end_time - start_time) / n_samples * 1000:.4f} ms")
```

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

### Basic Parallel Sampling with MLPModel

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/parallel_sampling.md

Initialize multiple chains and let TorchEBM handle parallelization for basic sampling. Requires importing torch, BaseModel, LangevinDynamics, and nn.

```python
import torch
from torchebm.core import BaseModel
from torchebm.samplers import LangevinDynamics
import torch.nn as nn

class MLPModel(BaseModel):
    def __init__(self, input_dim, hidden_dim=64):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, 1)
        )

    def forward(self, x):
        return self.network(x).squeeze(-1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MLPModel(input_dim=2, hidden_dim=32).to(device)

sampler = LangevinDynamics(
    model=model,
    step_size=0.1,
    noise_scale=0.01,
    device=device
)

n_samples = 10000
dim = 2
initial_points = torch.randn(n_samples, dim, device=device)

samples = sampler.sample(
    x=initial_points,
    n_steps=1000,
    return_trajectory=False
)

print(f"Generated {samples.shape[0]} samples of dimension {samples.shape[1]}")
```

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

### Batch Processing for Large Sample Sets

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/parallel_sampling.md

Process large numbers of samples in batches to avoid memory issues. This example iterates through batches, generating samples and accumulating them.

```python
import torch
import numpy as np
from torchebm.core import BaseModel
from torchebm.samplers import LangevinDynamics

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MLPModel(input_dim=2, hidden_dim=64).to(device)
sampler = LangevinDynamics(
    model=model,
    step_size=0.01,
    device=device
)

total_samples = 1000000
dim = 2
batch_size = 10000
num_batches = total_samples // batch_size

all_samples = np.zeros((total_samples, dim))

for i in range(num_batches):
    print(f"Generating batch {i+1}/{num_batches}")

    initial_points = torch.randn(batch_size, dim, device=device)

    batch_samples = sampler.sample(
        x=initial_points,
        n_steps=1000,
        return_trajectory=False
    )

    start_idx = i * batch_size
    end_idx = (i + 1) * batch_size
    all_samples[start_idx:end_idx] = batch_samples.cpu().numpy()

print(f"Generated {total_samples} samples in total")
```

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

### Use Langevin Dynamics Sampler

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/index.md

Demonstrates the usage of the `LangevinDynamics` sampler with a pre-defined energy function. This is a common pattern for generating samples from an energy model.

```python
from torchebm.samplers.langevin_dynamics import LangevinDynamics

sampler = LangevinDynamics(
    energy_function=energy_fn,
    step_size=0.01
)

samples = sampler.sample(
    dim=2, n_steps=100, n_samples=1000
)
```

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

### Basic HMC Sampling from GaussianModel

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/examples/samplers/index.md

Demonstrates how to initialize and use the HamiltonianMonteCarlo sampler with a GaussianModel. Ensure the model and initial particles are on the correct device.

```python
import torch
from torchebm.core import GaussianModel
from torchebm.samplers import HamiltonianMonteCarlo

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mean = torch.tensor([1.0, -1.0])
cov = torch.tensor([[1.0, 0.5], [0.5, 2.0]])
model = GaussianModel(mean, cov).to(device)

hmc_sampler = HamiltonianMonteCarlo(
    model=model,
    step_size=0.1,
    n_leapfrog_steps=10
)

initial_particles = torch.randn(1000, 2, device=device)
samples = hmc_sampler.sample(x=initial_particles, n_steps=25) # HMC often needs fewer steps

# Plotting is the same as before
```

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

### Equilibrium Matching Loss with Cosine Interpolant

Source: https://context7.com/soran-ghaderi/torchebm/llms.txt

Example of using EquilibriumMatchingLoss with a cosine interpolant for implicit energy matching. Requires model, prediction type, and energy type configuration.

```python
eqm = EquilibriumMatchingLoss(
    model=model,
    prediction="velocity",
    energy_type="none",        # implicit: model predicts gradient field
    interpolant="cosine",
    ct_threshold=0.8,
    ct_multiplier=4.0,
    apply_dispersion=False,
    device=device,
)

dataset = EightGaussiansDataset(n_samples=4000, seed=0)
loader  = DataLoader(dataset, batch_size=256, shuffle=True)
opt     = torch.optim.Adam(model.parameters(), lr=2e-4)

for epoch in range(5):
    for batch in loader:
        opt.zero_grad()
        loss = eqm(batch.to(device))
        loss.backward()
        opt.step()
    print(f"Epoch {epoch+1}  EqM loss: {loss.item():.4f}")
```

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

### Parameter Scheduling in Custom Samplers

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/samplers.md

Register and use schedulers to dynamically adjust sampler parameters during the sampling process. This example shows registering a constant scheduler and retrieving its value.

```python
self.register_scheduler("step_size", ConstantScheduler(0.01))

current_step_size = self.get_scheduled_value("step_size")

self.step_schedulers()
```

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

### Initialize EBM Training Components

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/examples/training/index.md

Initialize the MLP model, Langevin Dynamics sampler, ContrastiveDivergence loss function, and Adam optimizer. Configure sampler parameters and optimizer learning rate.

```python
# Model
model = MLPModel(input_dim=2).to(device)

# Sampler
sampler = LangevinDynamics(
    model=model,
    step_size=0.1,
    noise_scale=0.1,
)

# Loss Function
loss_fn = ContrastiveDivergence(
    model=model, 
    sampler=sampler, 
    n_steps=10 # k in CD-k
)

# Optimizer
optimizer = optim.Adam(model.parameters(), lr=1e-3)
```

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

### Generate a visual timeline

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/developer_guide/profiling.md

Use this command to generate a trace file (`trace.json`) for visualizing the execution timeline. Drag the `trace.json` file into [ui.perfetto.dev](https://ui.perfetto.dev) for analysis.

```bash
python benchmarks/profiler.py run --component FlowSampler --scale small --trace
```

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

### Plotting Synthetic 2D Datasets

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/examples/datasets/index.md

Generates and plots various synthetic 2D datasets available in TorchEBM. Imports necessary libraries and defines a plotting function. Ensure matplotlib is installed for visualization.

```python
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchebm.datasets import (
    GaussianMixtureDataset, EightGaussiansDataset, TwoMoonsDataset,
    SwissRollDataset, CircleDataset, CheckerboardDataset,
    PinwheelDataset, GridDataset
)

def plot_datasets(datasets, titles):
    fig, axes = plt.subplots(2, 4, figsize=(20, 10))
    axes = axes.flatten()
    for i, (data, title) in enumerate(zip(datasets, titles)):
        ax = axes[i]
        ax.scatter(data[:, 0], data[:, 1], s=5, alpha=0.7)
        ax.set_title(title)
        ax.grid(True, alpha=0.3)
        ax.axis('equal')
    plt.tight_layout()
    plt.show()

n_samples = 1000
seed = 42

datasets_to_plot = [
    GaussianMixtureDataset(n_samples=n_samples, n_components=8, std=0.07, radius=1.5, seed=seed).get_data(),
    EightGaussiansDataset(n_samples=n_samples, std=0.05, scale=2.0, seed=seed).get_data(),
    TwoMoonsDataset(n_samples=n_samples, noise=0.1, seed=seed).get_data(),
    SwissRollDataset(n_samples=n_samples, noise=0.1, arclength=3.0, seed=seed).get_data(),
    CircleDataset(n_samples=n_samples, noise=0.05, radius=1.0, seed=seed).get_data(),
    CheckerboardDataset(n_samples=n_samples, range_limit=4.0, noise=0.05, seed=seed).get_data(),
    PinwheelDataset(n_samples=n_samples, n_classes=5, noise=0.05, seed=seed).get_data(),
    GridDataset(n_samples_per_dim=30, range_limit=2.0, noise=0.02, seed=seed).get_data()
]
dataset_titles = [
    "Gaussian Mixture", "Eight Gaussians", "Two Moons", "Swiss Roll",
    "Circle", "Checkerboard", "Pinwheel", "2D Grid"
]

plot_datasets(datasets_to_plot, dataset_titles)
```

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

### ContrastiveDivergence Training with PCD

Source: https://context7.com/soran-ghaderi/torchebm/llms.txt

Uses Contrastive Divergence for CD-k training, approximating the log-likelihood gradient with MCMC steps. Supports energy regularization and persistent chains (PCD). Requires model, sampler, and dataset setup.

```python
import torch
from torchebm.core import BaseModel
from torchebm.samplers import LangevinDynamics
from torchebm.losses import ContrastiveDivergence
from torchebm.datasets import GaussianMixtureDataset
from torch.utils.data import DataLoader
import torch.nn as nn

class MLPEnergy(BaseModel):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2,64),nn.SiLU(),nn.Linear(64,64),nn.SiLU(),nn.Linear(64,1))
    def forward(self, x): return self.net(x).squeeze(-1)

device  = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model   = MLPEnergy().to(device)
sampler = LangevinDynamics(model=model, step_size=0.01, noise_scale=1.0, device=device)
cd_loss = ContrastiveDivergence(
    model=model, sampler=sampler,
    k_steps=20,              # number of MCMC steps
    energy_reg_weight=1e-3,  # regularize energy magnitude
    device=device,
)

dataset = GaussianMixtureDataset(n_samples=4000, n_components=4, seed=42)
loader  = DataLoader(dataset, batch_size=128, shuffle=True)
opt     = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):
    epoch_loss = 0.0
    for batch in loader:
        batch = batch.to(device)
        opt.zero_grad()
        loss, neg_samples = cd_loss(batch)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        epoch_loss += loss.item()
    print(f"Epoch {epoch+1:2d}  CD loss: {epoch_loss/len(loader):.4f}")
```

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

### Basic Langevin Dynamics Sampling

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/tutorials/samplers.md

Demonstrates the basic usage of Langevin Dynamics for sampling from a custom MLP model. Requires defining a model that inherits from BaseModel and initializing the sampler with model, step size, and noise scale.

```python
import torch
from torchebm.core import BaseModel
from torchebm.samplers import LangevinDynamics
import torch.nn as nn

class MLPModel(BaseModel):
    def __init__(self, input_dim, hidden_dim=64):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, 1)
        )

    def forward(self, x):
        return self.network(x).squeeze(-1)

model = MLPModel(input_dim=2, hidden_dim=32)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
langevin_sampler = LangevinDynamics(
    model=model,
    step_size=0.1,
    noise_scale=0.01,
    device=device
)

initial_points = torch.randn(100, 2, device=device)
samples = langevin_sampler.sample(
    x=initial_points,
    n_steps=1000,
    return_trajectory=False
)

print(samples.shape)
```

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

### Format Code and Run Tests

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/developer_guide/getting_started.md

Ensure your code adheres to formatting standards and all tests pass before committing. Use `black` and `isort` for formatting, and `pytest` for testing.

```bash
black torchebm/
    isort torchebm/
    pytest tests/ -v
```

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

### Langevin Dynamics Trajectory Visualization

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/examples/samplers/index.md

Visualizes the trajectories of Langevin Dynamics chains on a multimodal energy landscape. This example requires a custom BaseModel implementation and demonstrates how chains explore different modes. The plotting code is omitted for brevity.

```python
import numpy as np
import torch.nn as nn
from torchebm.core import BaseModel

class MultimodalModel(BaseModel):
    def __init__(self):
        super().__init__()
        self.centers = nn.Parameter(torch.tensor([
            [-1.5, -1.5], [1.5, 1.5], [-1.5, 1.5], [1.5, -1.5]
        ]), requires_grad=False)
        self.weights = nn.Parameter(torch.tensor([1.0, 1.0, 0.8, 0.8]), requires_grad=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dists = torch.cdist(x, self.centers)
        energy = -torch.logsumexp(-0.5 * dists.pow(2) * self.weights, dim=-1)
        return energy

model = MultimodalModel().to(device)
sampler = LangevinDynamics(model=model, step_size=0.1, noise_scale=0.1)

initial_particles = torch.zeros(5, 2, device=device) # 5 chains
trajectory = sampler.sample(x=initial_particles, n_steps=200, return_trajectory=True)

#  The plotting code (omitted for brevity, see full example for details)
# This would involve creating a contour plot of the energy landscape
# and overlaying the trajectories.
```

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

### Instantiate DoubleWellModel

Source: https://github.com/soran-ghaderi/torchebm/blob/master/docs/examples/models/index.md

Create a DoubleWellModel for testing sampler capabilities across energy barriers. Specify the barrier height.

```python
import torch
from torchebm.core import DoubleWellModel

double_well_model = DoubleWellModel(barrier_height=2.0)
```