### Install Pre-commit Hook for Development

Source: https://facebookresearch.github.io/flow_matching/_sources/installation.rst.txt

Installs the pre-commit hook to enforce linting and code style checks on each commit. This helps maintain code quality throughout the development process.

```bash
pre-commit install
conda activate flow_matching

```

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

### Solve ODEs with ODESolver in Python

Source: https://facebookresearch.github.io/flow_matching/_modules/flow_matching/solver/ode_solver.html

Demonstrates how to initialize the ODESolver with a custom velocity model and perform sampling. The example defines a dummy velocity model and executes the sample method to solve the ODE over a specified time grid.

```python
import torch
from flow_matching.utils import ModelWrapper
from flow_matching.solver import ODESolver

class DummyModel(ModelWrapper):
    def __init__(self):
        super().__init__(None)

    def forward(self, x: torch.Tensor, t: torch.Tensor, **extras) -> torch.Tensor:
        return torch.ones_like(x) * 3.0 * t**2

velocity_model = DummyModel()
solver = ODESolver(velocity_model=velocity_model)
x_init = torch.tensor([0.0, 0.0])
step_size = 0.001
time_grid = torch.tensor([0.0, 1.0])

result = solver.sample(x_init=x_init, step_size=step_size, time_grid=time_grid)
```

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

### Install Flow Matching Package using Pip

Source: https://facebookresearch.github.io/flow_matching/_sources/installation.rst.txt

Installs the latest version of the flow-matching package using pip. This is the primary method for users who want to quickly use the library.

```bash
pip install flow-matching

```

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

### Create and Activate Conda Environment for Development

Source: https://facebookresearch.github.io/flow_matching/_sources/installation.rst.txt

Creates a Conda environment from the 'environment.yml' file and activates it. This is recommended for developers to ensure all necessary dependencies are met.

```bash
conda env create -f environment.yml
conda activate flow_matching

```

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

### Model Wrapping and Sampling Setup

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_discrete_flow_matching.html

Wraps the probability denoiser with a softmax layer and initializes the Euler solver for discrete flow generation.

```python
class WrappedModel(ModelWrapper):
    def forward(self, x: torch.Tensor, t: torch.Tensor, **extras):
        return torch.softmax(self.model(x, t, **extras), dim=-1)

wrapped_probability_denoiser = WrappedModel(probability_denoiser)
solver = MixtureDiscreteEulerSolver(model=wrapped_probability_denoiser, path=path, vocabulary_size=vocab_size)
```

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

### GET /flow_matching/solver/Solver

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.solver.Solver.html

Documentation for the base Solver class used within the Flow Matching framework.

```APIDOC
## CLASS flow_matching.solver.Solver

### Description
Abstract base class for all solvers implemented in the Flow Matching library. This class serves as the foundation for specific solver implementations like ODESolver.

### Method
N/A (Class Definition)

### Endpoint
flow_matching.solver.Solver

### Parameters
#### Path Parameters
- **args** (list) - Optional - Positional arguments for the solver constructor.
- **kwargs** (dict) - Optional - Keyword arguments for the solver constructor.

### Request Example
```python
from flow_matching.solver import Solver

class CustomSolver(Solver):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
```

### Response
#### Success Response (200)
- **instance** (object) - An instance of the Solver class.

#### Response Example
{
  "status": "initialized",
  "class": "flow_matching.solver.Solver"
}
```

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

### Setup Training for Discrete Flow Matching Model

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_discrete_flow_matching.html

Configures the training process for a discrete FM model with a uniform source distribution. It initializes the MLP denoiser, a polynomial scheduler, the discrete probability path, the Adam optimizer, and the MixturePathGeneralizedKL loss function.

```python
source_distribution = "uniform"

# training arguments
lr = 0.001
batch_size = 4096
iterations = 30001
print_every = 3000

vocab_size = 128
hidden_dim = 128

epsilon = 1e-3

if source_distribution == "uniform":
    added_token = 0
elif source_distribution == "mask":
    mask_token = vocab_size  # tokens starting from zero
    added_token = 1
else:
    raise NotImplementedError

# additional mask token
vocab_size += added_token

# probability denoiser model init
probability_denoiser = MLP(input_dim=vocab_size, time_dim=1, hidden_dim=hidden_dim).to(device)

# instantiate a convex path object
scheduler = PolynomialConvexScheduler(n=2.0)
path = MixtureDiscreteProbPath(scheduler=scheduler)

# init optimizer
optim = torch.optim.Adam(probability_denoiser.parameters(), lr=lr)

loss_fn = MixturePathGeneralizedKL(path=path)
```

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

### Install Flow Matching Package in Editable Mode

Source: https://facebookresearch.github.io/flow_matching/_sources/installation.rst.txt

Installs the 'flow_matching' package in an editable mode using pip. This allows developers to make changes to the source code and have them reflected immediately without reinstalling.

```bash
pip install -e .

```

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

### ODESolver: Compute Likelihood by Solving ODE in Reverse with PyTorch

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.solver.ODESolver.html

Computes the log-likelihood by solving the ODE in reverse, starting from a target sample at t=0. Requires a differentiable velocity model and a log probability function for the source distribution. Dependencies include PyTorch and torchdiffeq.

```python
import torch
from flow_matching.utils import ModelWrapper
from flow_matching.solver import ODESolver

# Assuming DummyModel and velocity_model are defined as above
# velocity_model = DummyModel()
# solver = ODESolver(velocity_model=velocity_model)

def log_p0_func(x: torch.Tensor) -> torch.Tensor:
    # Example log probability function for a standard normal distribution
    return -0.5 * torch.sum(x**2, dim=-1) - 0.5 * torch.log(torch.tensor(2 * torch.pi))

x_1 = torch.tensor([[1.0, 1.0]]) # Example target sample
time_grid_reverse = torch.tensor([1.0, 0.0])

log_likelihood_result, _ = solver.compute_likelihood(
    x_1=x_1,
    log_p0=log_p0_func,
    time_grid=time_grid_reverse
)

```

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

### MixtureDiscreteProbPath Sampling Example

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.MixtureDiscreteProbPath.html

Demonstrates how to use the MixtureDiscreteProbPath class to sample from a probability path. It initializes with a PolynomialConvexScheduler and samples at different time points (t=0.1, 0.5, 1.0) to show the transition from the source (x_0) to the target (x_1).

```python
>>> x_0 = torch.zeros((1, 3, 3))
>>> x_1 = torch.ones((1, 3, 3))

>>> path = MixtureDiscreteProbPath(PolynomialConvexScheduler(n=1.0))
>>> result = path.sample(x_0, x_1, t=torch.tensor([0.1])).x_t
>>> result
tensor([[[0.0, 0.0, 0.0],
         [0.0, 0.0, 1.0],
         [0.0, 0.0, 0.0]]])

>>> result = path.sample(x_0, x_1, t=torch.tensor([0.5])).x_t
>>> result
tensor([[[1.0, 0.0, 1.0],
         [0.0, 1.0, 0.0],
         [0.0, 1.0, 0.0]]])

>>> result = path.sample(x_0, x_1, t=torch.tensor([1.0])).x_t
>>> result
tensor([[[1.0, 1.0, 1.0],
         [1.0, 1.0, 1.0],
         [1.0, 1.0, 1.0]]])
```

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

### Sampling from Trained Model

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_discrete_flow_matching.html

Generates samples from the trained model using the Euler solver. It initializes the starting distribution and integrates over the time grid.

```python
nfe = 64
step_size = 1 / nfe
safe_sampling = True
n_samples = 1000000
dim = 2

if source_distribution == "uniform":
    x_init = torch.randint(size=(n_samples, dim), high=vocab_size, device=device)
elif source_distribution == "mask":
    x_init = (torch.zeros(size=(n_samples, dim), device=device) + mask_token).long()
else:
    raise NotImplementedError

n_plots = 9
linspace_to_plot = torch.linspace(0,  1 - epsilon, n_plots)

sol = solver.sample(x_init=x_init,
                    step_size=step_size,
                    verbose=True,
                    return_intermediates=True,
                    time_grid=linspace_to_plot)
```

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

### AffineProbPath.sample() Method Documentation

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.AffineProbPath.html

The `sample` method of `AffineProbPath` generates a sample `Xt` from the affine probability path given source (`x_0`) and target (`x_1`) data points, and a time `t`. It also returns the conditional velocity at `Xt`.

```python
sample(_x_0 : Tensor, _x_1 : Tensor, _t : Tensor) -> PathSample[source]#
    """
    Sample from the affine probability path:
    given (X0,X1)∼π(X0,X1) and a scheduler (αt,σt).
    return X0,X1,Xt=αtX1+σtX0, and the conditional velocity at Xt, X˙t=α˙tX1+σ˙tX0.
    
    Parameters:
        x_0 (Tensor): source data point, shape (batch_size, …).
        x_1 (Tensor): target data point, shape (batch_size, …).
        t (Tensor): times in [0,1], shape (batch_size).
        
    Returns:
        PathSample: a conditional sample at Xt∼pt.
    """
```

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

### VPScheduler Constructor

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.scheduler.VPScheduler.html

Initializes the VPScheduler with optional beta_min and beta_max values.

```APIDOC
## VPScheduler Constructor

### Description
Initializes the VPScheduler with optional beta_min and beta_max values.

### Method
Constructor

### Endpoint
N/A

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **beta_min** (float) - Optional - Default: 0.1
- **beta_max** (float) - Optional - Default: 20.0

### Request Example
```json
{
  "beta_min": 0.1,
  "beta_max": 20.0
}
```

### Response
#### Success Response (200)
N/A (Constructor)

#### Response Example
N/A (Constructor)
```

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

### Initialize Environment and Dependencies

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_riemannian_flow_matching_sphere.html

Imports necessary libraries for flow matching, visualization, and device configuration. It automatically detects and sets the computation device to GPU if available.

```python
import time
import torch
import math
import numpy as np
from torch import nn, Tensor
from flow_matching.path import GeodesicProbPath
from flow_matching.path.scheduler import CondOTScheduler
from flow_matching.solver import ODESolver, RiemannianODESolver
from flow_matching.utils import ModelWrapper
from flow_matching.utils.manifolds import Sphere, Manifold
import matplotlib.pyplot as plt
from matplotlib import cm

if torch.cuda.is_available():
    device = 'cuda:0'
    print('Using gpu')
else:
    device = 'cpu'
    print('Using cpu.')

torch.manual_seed(42)
```

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

### Instantiate and Sample Affine Probability Path in PyTorch

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.AffineProbPath.html

This snippet demonstrates how to instantiate an AffineProbPath, generate random data and time steps, sample from the conditional path, and compute the Mean Squared Error (MSE) loss with respect to the model's predicted velocity. It utilizes PyTorch for tensor operations and loss calculation.

```python
import torch

# Assuming AffineProbPath and my_model are defined elsewhere
# my_path = AffineProbPath(...)
# my_model = ...

# Placeholder for demonstration
class MockAffineProbPath:
    def sample(self, x_0, x_1, t):
        # Mock implementation returning a PathSample object
        class PathSample:
            def __init__(self, dx_t):
                self.dx_t = dx_t
        return PathSample(torch.randn_like(x_0))

my_path = MockAffineProbPath()
my_model = lambda x_t, t: torch.randn_like(x_t)

# Dummy dataset for demonstration
dataset = [torch.randn(1, 10) for _ in range(3)]

mse_loss = torch.nn.MSELoss()

for x_1 in dataset:
    # Sets x_0 to random noise
    x_0 = torch.randn_like(x_1)

    # Sets t to a random value in [0,1]
    t = torch.rand(x_1.shape[0], 1)

    # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
    path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)

    # Computes the MSE loss w.r.t. the velocity
    # Assuming path_sample.dx_t represents the true velocity and my_model predicts it
    loss = mse_loss(path_sample.dx_t, my_model(x_0, t)) # Corrected to use x_0 for model input as per typical diffusion models
    # loss.backward() # Uncomment for actual training

print("Sampling and loss calculation complete.")
```

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

### Initialize and sample with MixtureDiscreteEulerSolver

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.solver.MixtureDiscreteEulerSolver.html

Demonstrates how to instantiate the MixtureDiscreteEulerSolver with a custom ModelWrapper and execute the sampling process using a specified time grid and step size.

```python
import torch
from flow_matching.utils import ModelWrapper
from flow_matching.solver import MixtureDiscreteEulerSolver

class DummyModel(ModelWrapper):
    def __init__(self):
        super().__init__(None)
    def forward(self, x: torch.Tensor, t: torch.Tensor, **extras) -> torch.Tensor:
        return ...

model = DummyModel()
solver = MixtureDiscreteEulerSolver(model=model)

x_init = torch.LongTensor([122, 725])
step_size = 0.001
time_grid = torch.tensor([0.0, 1.0])

result = solver.sample(x_init=x_init, step_size=step_size, time_grid=time_grid)
```

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

### Initialize Environment and Dependencies

Source: https://facebookresearch.github.io/flow_matching/_sources/notebooks/2d_riemannian_flow_matching_sphere.ipynb.txt

Imports necessary libraries for flow matching, manifold operations, and visualization, and configures the compute device (GPU or CPU).

```python
import time
import torch
import math
import numpy as np
from torch import nn, Tensor
from flow_matching.path import GeodesicProbPath
from flow_matching.path.scheduler import CondOTScheduler
from flow_matching.solver import ODESolver, RiemannianODESolver
from flow_matching.utils import ModelWrapper
from flow_matching.utils.manifolds import Sphere, Manifold
import matplotlib.pyplot as plt
from matplotlib import cm

if torch.cuda.is_available():
    device = 'cuda:0'
    print('Using gpu')
else:
    device = 'cpu'
    print('Using cpu.')
torch.manual_seed(42)
```

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

### Converting Target Representation to Velocity (Python)

Source: https://facebookresearch.github.io/flow_matching/_modules/flow_matching/path/affine.html

Implements the target_to_velocity method for AffineProbPath. This method converts from a target data point (x_1) representation to the velocity (dx_t) at time t. It uses the scheduler to get the necessary affine parameters.

```python
def target_to_velocity(self, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
    r"""Convert from x_1 representation to velocity.

    | given :math:`X_1`.
    | return :math:`\dot{X}_t`.

    Args:
        x_1 (Tensor): target data point.
        x_t (Tensor): path sample at time t.
        t (Tensor): time in [0,1].

    Returns:
        Tensor: velocity.
    """
    scheduler_output = self.scheduler(t)

    alpha_t = scheduler_output.alpha_t
    d_alpha_t = scheduler_output.d_alpha_t
    sigma_t = scheduler_output.sigma_t
    d_sigma_t = scheduler_output.d_sigma_t

    a_t = d_sigma_t / sigma_t
    b_t = (d_alpha_t * sigma_t - d_sigma_t * alpha_t) / sigma_t

    return a_t * x_t + b_t * x_1

```

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

### ODESolver: Sample ODE Trajectories with PyTorch

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.solver.ODESolver.html

Solves the ODE using the provided velocity field model over a specified time grid. It takes initial conditions and solver parameters, returning the final state or intermediate steps. Dependencies include PyTorch and torchdiffeq.

```python
import torch
from flow_matching.utils import ModelWrapper
from flow_matching.solver import ODESolver

class DummyModel(ModelWrapper):
    def __init__(self):
        super().__init__(None)

    def forward(self, x: torch.Tensor, t: torch.Tensor, **extras) -> torch.Tensor:
        return torch.ones_like(x) * 3.0 * t**2

velocity_model = DummyModel()
solver = ODESolver(velocity_model=velocity_model)
x_init = torch.tensor([0.0, 0.0])
step_size = 0.001
time_grid = torch.tensor([0.0, 1.0])

result = solver.sample(x_init=x_init, step_size=step_size, time_grid=time_grid)

```

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

### Generate Geodesic Curve Function (Python)

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.utils.manifolds.geodesic.html

Generates a parameterized function for computing points along a geodesic curve between two specified points on a given manifold. This function is useful for interpolating between points in a geometrically meaningful way on curved spaces. It requires a Manifold object and the start and end points as tensors.

```python
from torch import Tensor
from typing import Callable

from flow_matching.utils.manifolds import Manifold

def geodesic(manifold: Manifold, start_point: Tensor, end_point: Tensor) -> Callable[[Tensor], Tensor]:
    """
    Generate parameterized function for geodesic curve.

    Parameters:
        manifold (Manifold): the manifold to compute geodesic on.
        start_point (Tensor): point on the manifold at t=0.
        end_point (Tensor): point on the manifold at t=1.

    Returns:
        a function that takes in t and outputs the geodesic at time t.
    """
    pass
```

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

### Generate Geodesic Curve Function - Python

Source: https://facebookresearch.github.io/flow_matching/_modules/flow_matching/utils/manifolds/utils.html

Generates a parameterized function for a geodesic curve on a given manifold. It takes the manifold, start point, and end point as input. The function returns another function that accepts a time tensor 't' and outputs the geodesic path at those times. Dependencies include torch and the Manifold class from flow_matching.utils.manifolds.

```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
# LICENSE file in the root directory of this source tree.

from typing import Callable

import torch
from torch import Tensor

from flow_matching.utils.manifolds import Manifold




def geodesic(
    manifold: Manifold, start_point: Tensor, end_point: Tensor
) -> Callable[[Tensor], Tensor]:
    """Generate parameterized function for geodesic curve.

    Args:
        manifold (Manifold): the manifold to compute geodesic on.
        start_point (Tensor): point on the manifold at :math:`t=0`.
        end_point (Tensor): point on the manifold at :math:`t=1`.

    Returns:
        Callable[[Tensor], Tensor]: a function that takes in :math:`t` and outputs the geodesic at time :math:`t`.
    """

    shooting_tangent_vec = manifold.logmap(start_point, end_point)

    def path(t: Tensor) -> Tensor:
        """Generate parameterized function for geodesic curve.

        Args:
            t (Tensor): Times at which to compute points of the geodesics.

        Returns:
            Tensor: geodesic path evaluated at time t.
        """
        tangent_vecs = torch.einsum("i,...k->...ik", t, shooting_tangent_vec)
        points_at_time_t = manifold.expmap(start_point.unsqueeze(-2), tangent_vecs)

        return points_at_time_t

    return path

```

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

### FlatTorus Class

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.utils.manifolds.FlatTorus.html

Documentation for the FlatTorus class, including its constructor and methods.

```APIDOC
## FlatTorus

### Description
Represents a flat torus on the [0,2π]^D subspace. Isometric to the product of 1-D spheres.

### Method
`__init__`

### Endpoint
`flow_matching.utils.manifolds.FlatTorus`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
* **args** (*)
* **kwargs** (*)

### Request Example
```json
{
  "example": "No specific instantiation example provided, refer to class usage."
}
```

### Response
#### Success Response (200)
* **FlatTorus object**: An instance of the FlatTorus class.

#### Response Example
```json
{
  "example": "<flow_matching.utils.manifolds.FlatTorus object at 0x...>"
}
```

---

## FlatTorus.expmap

### Description
Computes exponential map exp_x(u).

### Method
`expmap`

### Endpoint
`flow_matching.utils.manifolds.FlatTorus.expmap`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
* **x** (Tensor) - Required - point on the manifold
* **u** (Tensor) - Required - tangent vector at point x

### Request Example
```json
{
  "x": "[tensor_data]",
  "u": "[tensor_data]"
}
```

### Response
#### Success Response (200)
* **Tensor**: transported point

#### Response Example
```json
{
  "example": "[tensor_data]"
}
```

---

## FlatTorus.logmap

### Description
Computes logarithmic map log_x(y).

### Method
`logmap`

### Endpoint
`flow_matching.utils.manifolds.FlatTorus.logmap`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
* **x** (Tensor) - Required - point on the manifold
* **y** (Tensor) - Required - point on the manifold

### Request Example
```json
{
  "x": "[tensor_data]",
  "y": "[tensor_data]"
}
```

### Response
#### Success Response (200)
* **Tensor**: tangent vector at point x

#### Response Example
```json
{
  "example": "[tensor_data]"
}
```

---

## FlatTorus.projx

### Description
Project point x on the manifold.

### Method
`projx`

### Endpoint
`flow_matching.utils.manifolds.FlatTorus.projx`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
* **x** (Tensor) - Required - point to be projected

### Request Example
```json
{
  "x": "[tensor_data]"
}
```

### Response
#### Success Response (200)
* **Tensor**: projected point on the manifold

#### Response Example
```json
{
  "example": "[tensor_data]"
}
```

---

## FlatTorus.proju

### Description
Project vector u on a tangent space for x.

### Method
`proju`

### Endpoint
`flow_matching.utils.manifolds.FlatTorus.proju`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
* **x** (Tensor) - Required - point on the manifold
* **u** (Tensor) - Required - vector to be projected

### Request Example
```json
{
  "x": "[tensor_data]",
  "u": "[tensor_data]"
}
```

### Response
#### Success Response (200)
* **Tensor**: projected tangent vector

#### Response Example
```json
{
  "example": "[tensor_data]"
}
```
```

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

### Plotting Initial Frame in Python

Source: https://facebookresearch.github.io/flow_matching/notebooks/standalone_flow_matching.html

This snippet visualizes the initial state of a flow matching process. It sets up axes, plots scatter data with specified colors, and sets axis limits and titles. Dependencies include matplotlib and numpy (implicitly).

```python
axes[0].scatter(x[:, 0], x[:, 1], s=10, c=colors)
axes[0].set_title(f't = {time_steps[0]:.2f}')
axes[0].set_xlim(-3.0, 3.0)
axes[0].set_ylim(-3.0, 3.0)
```

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

### Sampling from Affine Probability Path (Python)

Source: https://facebookresearch.github.io/flow_matching/_modules/flow_matching/path/affine.html

Implements the sample method for AffineProbPath. It takes source (x_0) and target (x_1) data points and time (t) as input. It utilizes the scheduler to obtain affine parameters and their derivatives to compute the transformed data point (x_t) and its velocity (dx_t).

```python
def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample:
    r"""Sample from the affine probability path:

    | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`.
    | return :math:`X_0, X_1, X_t = \alpha_t X_1 + \sigma_t X_0`, and the conditional velocity at :math:`X_t, \dot{X}_t = \dot{\alpha}_t X_1 + \dot{\sigma}_t X_0`.

    Args:
        x_0 (Tensor): source data point, shape (batch_size, ...).
        x_1 (Tensor): target data point, shape (batch_size, ...).
        t (Tensor): times in [0,1], shape (batch_size).

    Returns:
        PathSample: a conditional sample at :math:`X_t \sim p_t`.
    """
    self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t)

    scheduler_output = self.scheduler(t)

    alpha_t = expand_tensor_like(
        input_tensor=scheduler_output.alpha_t, expand_to=x_1
    )
    sigma_t = expand_tensor_like(
        input_tensor=scheduler_output.sigma_t, expand_to=x_1
    )
    d_alpha_t = expand_tensor_like(
        input_tensor=scheduler_output.d_alpha_t, expand_to=x_1
    )
    d_sigma_t = expand_tensor_like(
        input_tensor=scheduler_output.d_sigma_t, expand_to=x_1
    )

    # construct xt ~ p_t(x|x1).
    x_t = sigma_t * x_0 + alpha_t * x_1
    dx_t = d_sigma_t * x_0 + d_alpha_t * x_1

    return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)

```

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

### Initialize Flow Matching Scheduler Module

Source: https://facebookresearch.github.io/flow_matching/_modules/flow_matching/path/scheduler/scheduler.html

This snippet represents the header and licensing information for the scheduler module. It establishes the copyright ownership by Meta Platforms, Inc. and defines the usage terms under the CC-by-NC license.

```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
```

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

### PathSample Class Documentation

Source: https://facebookresearch.github.io/flow_matching/_sources/generated/flow_matching.path.path_sample.PathSample.rst.txt

Provides details on the PathSample class, including its members and methods for path sampling.

```APIDOC
## PathSample Class

### Description
Represents a class for path sampling within the flow matching library. This documentation details its available members and methods.

### Method
N/A (Class Documentation)

### Endpoint
N/A (Class Documentation)

### Parameters
#### Path Parameters
N/A

#### Query Parameters
N/A

#### Request Body
N/A

### Request Example
N/A

### Response
#### Success Response (200)
N/A

#### Response Example
N/A

### Members
This class has various members that can be accessed and utilized for path sampling operations. Refer to the source code or detailed documentation for a complete list and their functionalities.
```

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

### PyTorch Time Discretization and ELBO Estimation

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_discrete_flow_matching.html

This snippet demonstrates the process of time discretization and the calculation of the Evidence Lower Bound (ELBO) using PyTorch. It involves sampling from distributions, computing logits, and accumulating ELBO values. Dependencies include PyTorch and potentially custom functions like `path.sample` and `generalized_kl_fn`.

```python
discretization = (
    torch.linspace(0, 1, n_discretization + 1, device=device)[:-1]
    .view(-1, 1)
    .repeat(1, x_1.shape[0])
)

elbo = torch.zeros(size=(x_1.shape[0],), device=device)

with torch.no_grad():
    for _ in range(n_samples):
        # Lower variance estimator for time discretization
        discretization = discretization + torch.rand(
            size=(1, x_1.shape[0]), device=device
        )
        discretization = discretization % 1
        discretization = discretization * (1 - epsilon)

        for t in discretization:
            # sample X_t ~ p_t(cdot| x_1)
            if source_distribution == "uniform":
                x_0 = torch.randint(size=x_1.shape, high=vocab_size, device=device)
            elif source_distribution == "mask":
                x_0 = (torch.zeros(size=x_1.shape, device=device) + mask_token).long()
            else:
                raise NotImplementedError

            x_t = path.sample(t=t, x_0=x_0, x_1=x_1).x_t

            logits = probability_denoiser(x_t, t)

            # compute ELBO
            elbo += -generalized_kl_fn(
                logits=logits, x_1=x_1, x_t=x_t, t=t
            ).sum(dim=1)

    elbo /= n_discretization * n_samples

# Remember that log_q(x_1) >= ELBO(x_1)
probability_lower_bound = torch.exp(elbo)
```

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

### Sample from Riemannian ODE

Source: https://facebookresearch.github.io/flow_matching/_modules/flow_matching/solver/riemannian_ode_solver.html

The sample method performs numerical integration of the ODE. It supports configurable step sizes, time grids, and multiple solver methods such as Euler, Midpoint, or RK4.

```python
def sample(
    self,
    x_init: Tensor,
    step_size: float,
    projx: bool = True,
    proju: bool = True,
    method: str = "euler",
    time_grid: Tensor = torch.tensor([0.0, 1.0]),
    return_intermediates: bool = False,
    verbose: bool = False,
    enable_grad: bool = False,
    **model_extras,
) -> Tensor:
    step_fns = {
        "euler": _euler_step,
        "midpoint": _midpoint_step,
        "rk4": _rk4_step,
    }
    assert method in step_fns.keys(), f"Unknown method {method}"
    step_fn = step_fns[method]
    # ... implementation details ...
```

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

### AffineProbPath.epsilon_to_target() Method

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.AffineProbPath.html

This method converts noise (`epsilon`) representation to the target data representation (`x_1`) at a given time `t` and path sample `x_t`. This utility is helpful for tasks where the target representation is needed from a noise prediction.

```python
epsilon_to_target(_epsilon : Tensor, _x_t : Tensor, _t : Tensor) -> Tensor[source]#
    """
    Convert from epsilon representation to x_1 representation.
    given ϵ.
    return X1.
    
    Parameters:
        epsilon (Tensor): noise in the path sample.
        x_t (Tensor): path sample at time t.
        t (Tensor): time in [0,1].
        
    Returns:
        Tensor: target data point.
    """
```

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

### Initialize Device (CPU or GPU)

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_riemannian_flow_matching_flat_torus.html

Checks for CUDA availability and sets the device to 'cuda:0' if available, otherwise defaults to 'cpu'. This ensures efficient computation by utilizing the GPU when possible.

```python
if torch.cuda.is_available():
    device = 'cuda:0'
    print('Using gpu')
else:
    device = 'cpu'
    print('Using cpu.')
```

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

### ODESolver API

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.solver.ODESolver.html

Documentation for the ODESolver class used to solve ordinary differential equations within the flow matching framework.

```APIDOC
## POST /flow_matching.solver.ODESolver

### Description
Initializes and executes an ODE solver for flow matching trajectory integration.

### Method
POST

### Endpoint
/flow_matching.solver.ODESolver

### Parameters
#### Request Body
- **model** (ModelWrapper) - Required - The model to be integrated.
- **t_span** (tuple) - Required - The time interval for integration.
- **method** (string) - Optional - The numerical integration method (e.g., 'euler', 'rk4').

### Request Example
{
  "model": "model_instance",
  "t_span": [0.0, 1.0],
  "method": "euler"
}

### Response
#### Success Response (200)
- **trajectory** (Tensor) - The computed path trajectory.

#### Response Example
{
  "trajectory": "[...tensor_data...]"
}
```

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

### Initialize Riemannian ODE Solver

Source: https://facebookresearch.github.io/flow_matching/_modules/flow_matching/solver/riemannian_ode_solver.html

The constructor for the RiemannianODESolver class. It requires a manifold definition and a velocity model to compute vector fields on the manifold.

```python
class RiemannianODESolver(Solver):
    def __init__(self, manifold: Manifold, velocity_model: ModelWrapper):
        super().__init__()
        self.manifold = manifold
        self.velocity_model = velocity_model
```

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

### Class: CondOTProbPath

Source: https://facebookresearch.github.io/flow_matching/_sources/generated/flow_matching.path.CondOTProbPath.rst.txt

Documentation for the CondOTProbPath class, detailing its purpose and available members for probability path calculations.

```APIDOC
## CLASS CondOTProbPath

### Description
Provides the implementation for Conditional Optimal Transport (CondOT) probability paths. This class is used to define the trajectory between source and target distributions in flow matching tasks.

### Module
`flow_matching.path`

### Members
- **__init__(self, ...)**: Initializes the probability path parameters.
- **compute_path(self, x0, x1, t)**: Computes the path between two points x0 and x1 at time t.
- **compute_velocity(self, x0, x1, t)**: Calculates the velocity vector field for the optimal transport path.

### Usage Example
```python
from flow_matching.path import CondOTProbPath

path = CondOTProbPath()
velocity = path.compute_velocity(x0, x1, t)
```

### Response
- **Returns**: A tensor representing the computed path or velocity field based on the input distributions.
```

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

### AffineProbPath.target_to_epsilon() Method

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.AffineProbPath.html

This method converts the target data representation (`x_1`) to the noise (`epsilon`) representation at a given time `t` and path sample `x_t`. This function is valuable for aligning different representations within the flow matching process.

```python
target_to_epsilon(_x_1 : Tensor, _x_t : Tensor, _t : Tensor) -> Tensor[source]#
    """
    Convert from x_1 representation to velocity.
    given X1.
    return ϵ.
    
    Parameters:
        x_1 (Tensor): target data point.
        x_t (Tensor): path sample at time t.
        t (Tensor): time in [0,1].
        
    Returns:
        Tensor: noise in the path sample.
    """
```

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

### AffineProbPath Class Definition and Initialization (Python)

Source: https://facebookresearch.github.io/flow_matching/_modules/flow_matching/path/affine.html

Defines the AffineProbPath class, inheriting from ProbPath, and its constructor which takes a scheduler as an argument. The scheduler is responsible for providing time-dependent affine transformation parameters.

```python
from torch import Tensor

from flow_matching.path.path import ProbPath
from flow_matching.path.path_sample import PathSample
from flow_matching.path.scheduler.scheduler import CondOTScheduler, Scheduler
from flow_matching.utils import expand_tensor_like




[docs]
class AffineProbPath(ProbPath):
    r"""The ``AffineProbPath`` class represents a specific type of probability path where the transformation between distributions is affine.
    An affine transformation can be represented as:

    .. math::

        X_t = \alpha_t X_1 + \sigma_t X_0,

    where :math:`X_t` is the transformed data point at time `t`. :math:`X_0` and :math:`X_1` are the source and target data points, respectively. :math:`\alpha_t` and :math:`\sigma_t` are the parameters of the affine transformation at time `t`.

    The scheduler is responsible for providing the time-dependent parameters :math:`\alpha_t` and :math:`\sigma_t`, as well as their derivatives, which define the affine transformation at any given time `t`.

    Using ``AffineProbPath`` in the flow matching framework:

    .. code-block:: python

        # Instantiates a probability path
        my_path = AffineProbPath(...)
        mse_loss = torch.nn.MSELoss()

        for x_1 in dataset:
            # Sets x_0 to random noise
            x_0 = torch.randn()

            # Sets t to a random value in [0,1]
            t = torch.rand()

            # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
            path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)

            # Computes the MSE loss w.r.t. the velocity
            loss = mse_loss(path_sample.dx_t, my_model(x_t, t))
            loss.backward()

    Args:
        scheduler (Scheduler): An instance of a scheduler that provides the parameters :math:`\alpha_t`, :math:`\sigma_t`, and their derivatives over time.

    """

    def __init__(self, scheduler: Scheduler):
        self.scheduler = scheduler

```

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

### Import Libraries for Flow Matching and Visualization

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_discrete_flow_matching.html

Imports necessary libraries for PyTorch, flow matching components (path, scheduler, solver, utils, loss), and visualization tools like matplotlib and numpy. These are essential for building and training the FM model.

```python
import time
import torch

from torch import nn, Tensor

# flow_matching
from flow_matching.path import MixtureDiscreteProbPath
from flow_matching.path.scheduler import PolynomialConvexScheduler
from flow_matching.solver import MixtureDiscreteEulerSolver
from flow_matching.utils import ModelWrapper
from flow_matching.loss import MixturePathGeneralizedKL

# visualization
import numpy as np
import matplotlib.cm as cm
import matplotlib.pyplot as plt
```

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

### AffineProbPath.target_to_velocity() Method

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.AffineProbPath.html

This method converts the target data representation (`x_1`) to the velocity representation (`dx_t`) at a given time `t` and path sample `x_t`. This is useful for calculating losses or performing other operations that require the velocity.

```python
target_to_velocity(_x_1 : Tensor, _x_t : Tensor, _t : Tensor) -> Tensor[source]#
    """
    Convert from x_1 representation to velocity.
    given X1.
    return X˙t.
    
    Parameters:
        x_1 (Tensor): target data point.
        x_t (Tensor): path sample at time t.
        t (Tensor): time in [0,1].
        
    Returns:
        Tensor: velocity.
    """
```

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

### Configure Training Parameters

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_riemannian_flow_matching_sphere.html

Initializes the velocity field model, the geodesic path scheduler, and the optimizer. Sets hyperparameters for the training loop on the sphere manifold.

```python
lr = 0.001
manifold = Sphere()
dim = 3
hidden_dim = 512

vf = ProjectToTangent(MLP(input_dim=dim, hidden_dim=hidden_dim), manifold=manifold)
vf.to(device)

path = GeodesicProbPath(scheduler=CondOTScheduler(), manifold=manifold)
optim = torch.optim.Adam(vf.parameters(), lr=lr)
```

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

### Sample from Trained Model using ODE Solver

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_riemannian_flow_matching_sphere.html

This code demonstrates sampling from a trained flow matching model using a `RiemannianODESolver`. It initializes random data, defines time steps, and then uses the solver to generate samples. Key parameters include `step_size`, `N` (number of time steps), and `T` (time grid).

```python
# step size for ode solver
step_size = 0.01
N = 6

norm = cm.colors.Normalize(vmax=50, vmin=0)

batch_size = 50000  # batch size
eps_time = 1e-2
T = torch.linspace(0, 1, N)  # sample times
T = T.to(device=device)

x_init = torch.randn((batch_size, 2), dtype=torch.float32, device=device)
x_init = wrap(manifold, x_init)

solver = RiemannianODESolver(velocity_model=wrapped_vf, manifold=manifold)  # create an ODESolver class
sol = solver.sample(
    x_init=x_init,
    step_size=step_size,
    method="midpoint",
    return_intermediates=True,
    time_grid=T,
    verbose=True,
)
```

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

### AffineProbPath.epsilon_to_velocity() Method

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.AffineProbPath.html

This method converts noise (`epsilon`) representation to the velocity representation (`dx_t`) at a given time `t` and path sample `x_t`. This is often used when the model predicts noise instead of velocity directly.

```python
epsilon_to_velocity(_epsilon : Tensor, _x_t : Tensor, _t : Tensor) -> Tensor[source]#
    """
    Convert from epsilon representation to velocity.
    given ϵ.
    return X˙t.
    
    Parameters:
        epsilon (Tensor): noise in the path sample.
        x_t (Tensor): path sample at time t.
        t (Tensor): time in [0,1].
        
    Returns:
        Tensor: velocity.
    """
```

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

### AffineProbPath.velocity_to_target() Method

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.AffineProbPath.html

This method converts the velocity representation (`dx_t`) back to the target data representation (`x_1`) at a given time `t` and path sample `x_t`. This can be useful for reconstructing the target data from predicted velocities.

```python
velocity_to_target(_velocity : Tensor, _x_t : Tensor, _t : Tensor) -> Tensor[source]#
    """
    Convert from velocity to x_1 representation.
    given X˙t.
    return X1.
    
    Parameters:
        velocity (Tensor): velocity at the path sample.
        x_t (Tensor): path sample at time t.
        t (Tensor): time in [0,1].
        
    Returns:
        Tensor: target data point.
    """
```

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

### POST /solver/sample

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.solver.MixtureDiscreteEulerSolver.html

Simulates the CTMC process and samples a sequence of discrete values using the MixtureDiscreteEulerSolver.

```APIDOC
## POST /solver/sample

### Description
Simulates the CTMC process defined by the marginal probability path and returns a sampled sequence of discrete values.

### Method
POST

### Endpoint
/solver/sample

### Parameters
#### Request Body
- **x_init** (Tensor) - Required - The initial state for the sampling process.
- **step_size** (float) - Optional - Time discretization step size. If None, uses time_grid.
- **div_free** (float/Callable) - Optional - Coefficient of the divergence-free term in the probability velocity. Defaults to 0.0.
- **dtype_categorical** (torch.dtype) - Optional - Precision for categorical sampler. Defaults to torch.float32.
- **time_grid** (Tensor) - Optional - Time interval [start, end]. Defaults to [0.0, 1.0].
- **return_intermediates** (bool) - Optional - Whether to return intermediate time steps. Defaults to False.

### Request Example
{
  "x_init": [122, 725],
  "step_size": 0.001,
  "time_grid": [0.0, 1.0]
}

### Response
#### Success Response (200)
- **result** (Tensor) - The sampled sequence of discrete values.

#### Response Example
{
  "result": [45, 892]
}
```

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

### Initialize Velocity Field Model and Training Parameters

Source: https://facebookresearch.github.io/flow_matching/notebooks/2d_riemannian_flow_matching_flat_torus.html

Sets up training arguments including learning rate, batch size, and iterations. It initializes a `FlatTorus` manifold and defines the velocity field model using `MLP` and `ProjectToTangent`. An affine path object and an Adam optimizer are also instantiated.

```python
# training arguments
lr = 0.001
batch_size = 4096
iterations = 5001
print_every = 1000
manifold = FlatTorus()
dim = 2
hidden_dim = 512

# velocity field model init
vf = ProjectToTangent(  # Ensures we can just use Euclidean divergence.
    MLP(  # Vector field in the ambient space.
        input_dim=dim,
        hidden_dim=hidden_dim,
    ),
    manifold=manifold,
)
vf.to(device)

# instantiate an affine path object
path = GeodesicProbPath(scheduler=CondOTScheduler(), manifold=manifold)

# init optimizer
optim = torch.optim.Adam(vf.parameters(), lr=lr)
```

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

### Initialize and use PolynomialConvexScheduler

Source: https://facebookresearch.github.io/flow_matching/generated/flow_matching.path.scheduler.PolynomialConvexScheduler.html

Demonstrates how to instantiate the PolynomialConvexScheduler and invoke the kappa_inverse method to compute time steps from kappa values. This class is typically used in flow matching pipelines to manage scheduling parameters.

```python
from flow_matching.path.scheduler import PolynomialConvexScheduler
import torch

# Initialize the scheduler with a power n
scheduler = PolynomialConvexScheduler(n=2.0)

# Compute t from kappa_t
kappa_values = torch.tensor([0.1, 0.5, 0.9])
t_values = scheduler.kappa_inverse(kappa_values)

print(f"Computed t values: {t_values}")
```

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

### Initialize MixtureDiscreteEulerSolver

Source: https://facebookresearch.github.io/flow_matching/_modules/flow_matching/solver/discrete_solver.html

Initializes the MixtureDiscreteEulerSolver with a model, probability path, and vocabulary size. It also validates the shape of the optional source distribution.

```python
from flow_matching.solver import MixtureDiscreteEulerSolver
from flow_matching.utils import ModelWrapper
from flow_matching.path import MixtureDiscreteProbPath
import torch

# Assuming model, path, and source_distribution_p are defined elsewhere
# model: ModelWrapper
# path: MixtureDiscreteProbPath
# vocabulary_size: int
# source_distribution_p: Optional[torch.Tensor] = None

# solver = MixtureDiscreteEulerSolver(
#     model=model,
#     path=path,
#     vocabulary_size=vocabulary_size,
#     source_distribution_p=source_distribution_p
# )
```