### Install TorchDEQ

Source: https://github.com/locuslab/torchdeq/blob/main/README.md

Installation instructions for TorchDEQ using pip or from source.

```bash
pip install torchdeq
```

```bash
git clone https://github.com/locuslab/torchdeq.git && cd torchdeq
pip install -e .
```

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

### DEQ Transformer Training and Inference Example

Source: https://context7.com/locuslab/torchdeq/llms.txt

Provides a complete example of setting up the DEQ Transformer for training and inference. It includes argument parsing for DEQ parameters, model instantiation, optimizer setup, a single training step, and an inference pass.

```python
# Training example
parser = argparse.ArgumentParser()
add_deq_args(parser)
args = parser.parse_args([
    '--f_solver', 'anderson',
    '--f_max_iter', '24',
    '--n_states', '2',
    '--gamma', '0.8',
    '--jac_loss_weight', '0.1',
    '--norm_type', 'weight_norm'
])

model = DEQTransformer(args, vocab_size=10000)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# Training step
model.train()
batch_ids = torch.randint(0, 10000, (4, 64))
targets = torch.randint(0, 10000, (4, 64))

loss, info = model(batch_ids, targets)
loss.backward()
optimizer.step()

print(f"Loss: {loss.item():.4f}")
print(f"Convergence: {info['abs_lowest'].mean().item():.2e}")
print(f"Steps: {info['nstep'].mean().item():.1f}")

# Inference
model.eval()
with torch.no_grad():
    logits, info = model(batch_ids)
    predictions = logits.argmax(dim=-1)
    print(f"Inference convergence: {info['abs_lowest'].mean().item():.2e}")
```

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

### Install TorchDEQ via Package Managers or Source

Source: https://github.com/locuslab/torchdeq/blob/main/docs/get_started.md

Commands to install the TorchDEQ library using pip, conda, or by cloning the repository from source.

```bash
pip install torchdeq
conda install torchdeq
git clone https://github.com/locuslab/torchdeq.git && cd torchdeq
pip install -e .
```

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

### Install Project Dependencies

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/ignn/README.md

Installs the required Python packages listed in the requirements.txt file to ensure the environment is correctly configured for IGNN.

```bash
pip install -r requirements.txt
```

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

### Full DEQ Model Implementation Example

Source: https://github.com/locuslab/torchdeq/blob/main/docs/get_started.md

A complete example showing the integration of DEQ into a PyTorch module, including normalization, forward pass logic, and a training loop.

```python
import argparse
import torch
from torchdeq import get_deq, apply_norm, reset_norm
from torchdeq.utils import add_deq_args

class DEQDemo(torch.nn.Module):
    def __init__(self, args):
        super().__init__()
        self.deq_func = DEQFunc(args)
        apply_norm(self.deq_func, args=args)
        self.deq = get_deq(args)

    def forward(self, x, z0):
        reset_norm(self.deq_func)
        f = lambda z: self.deq_func(z, x)
        return self.deq(f, z0)

def train(args, inj, deq, decoder, loader, loss, opt):
    for x, y in loader:
        z0 = torch.randn(args.z_shape)
        z_out, info = deq(inj(x), z0)
        l = loss(decoder(z_out[-1]), y)
        l.backward()
        opt.step()
```

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

### Install Dependencies with Conda

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-flow/README.md

Installs the necessary dependencies for the project using Conda. This includes Python, PyTorch, torchvision, cudatoolkit, and other scientific libraries. It's recommended to create a dedicated Conda environment for this project.

```bash
conda create --name deq python==3.7.9
conda activate deq
conda install pytorch==1.11.0 torchvision==0.12.0 cudatoolkit=11.3 -c pytorch -c conda-forge
conda install tensorboard scipy opencv matplotlib einops -c conda-forge
```

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

### Register and Use Custom Solvers

Source: https://context7.com/locuslab/torchdeq/llms.txt

Allows registration of custom fixed-point solvers with the DEQ framework using `register_solver`. The example demonstrates registering a Newton-like solver and then retrieving and using it via `get_solver`.

```python
import torch
from torchdeq.solver import register_solver, get_solver
from torchdeq.solver.utils import solver_stat_from_final_step

def newton_solver(func, x0, max_iter=20, tol=1e-6, **kwargs):
    """Simple Newton-like solver for demonstration."""
    x = x0
    for i in range(max_iter):
        with torch.enable_grad():
            x_input = x.clone().requires_grad_(True)
            fx = func(x_input)
            gx = fx - x_input

            # Approximate Newton step
            x = fx.detach() # This is a simplification, a true Newton step would involve Jacobians

            if gx.abs().max() < tol:
                break

    info = solver_stat_from_final_step(x0, x, i + 1)
    return x, [], info

# Register the custom solver
register_solver('newton', newton_solver)

# Use the custom solver
solver = get_solver('newton')
f = lambda z: 0.5 * (z + 2.0 / z)
z_star, _, info = solver(f, torch.tensor([[1.5]]))
print(f"Newton solver result: {z_star.item():.6f}")
```

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

### Install NVIDIA DALI for ImageNet Training

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/mdeq/README.md

Installs the NVIDIA DALI library, which is recommended for faster ImageNet training. This command specifies a particular version of DALI compatible with CUDA 10.0.

```bash
pip3 install --extra-index-url https://developer.download.nvidia.com/compute/redist/cuda/10.0 nvidia-dali==0.13.0
```

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

### Install Dependencies with Conda

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-ddim/README.md

Installs project dependencies using Conda by creating a new environment from a requirements file. This ensures a consistent and reproducible development environment.

```bash
conda create --name [ENV_NAME] --file requirements.txt
conda activate [ENV_NAME]
```

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

### Create DEQ Module with get_deq

Source: https://context7.com/locuslab/torchdeq/llms.txt

The get_deq factory function generates a DEQ module, allowing customization of computational graphs, solvers, iteration limits, tolerances, and differentiation methods. It can be configured directly or via argparse. The example demonstrates creating a DEQ, defining a fixed-point function, and solving for the equilibrium.

```python
import torch
import torch.nn as nn
from torchdeq import get_deq, reset_deq
from torchdeq.norm import apply_norm

# Create DEQ with default settings (sliced computational graph)
deq = get_deq()

# Create DEQ with specific configuration
deq = get_deq(
    core='indexing',           # 'indexing' or 'sliced' computational graph
    f_solver='anderson',       # Forward solver: 'anderson', 'broyden', 'fixed_point_iter'
    b_solver='fixed_point_iter',  # Backward solver
    f_max_iter=40,             # Max forward iterations
    b_max_iter=40,             # Max backward iterations
    f_tol=1e-3,                # Forward tolerance
    b_tol=1e-6,                # Backward tolerance
    f_stop_mode='abs',         # Stop mode: 'abs' or 'rel'
    ift=True,                  # Enable Implicit Function Theorem
    n_states=3,                # Number of trajectory states to sample
    grad=1,                    # PhantomGrad steps
    tau=1.0                    # Damping factor
)

# Using argparse configuration
import argparse
from torchdeq import add_deq_args

parser = argparse.ArgumentParser()
add_deq_args(parser)
args = parser.parse_args([])
deq = get_deq(args)

# Define a simple fixed-point function
class DEQLayer(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.linear = nn.Linear(hidden_dim, hidden_dim)
        self.activation = nn.ReLU()

    def forward(self, z, x):
        return self.activation(self.linear(z) + x)

# Forward pass through DEQ
hidden_dim, batch_size = 64, 16
layer = DEQLayer(hidden_dim)
apply_norm(layer)  # Apply weight normalization

x = torch.randn(batch_size, hidden_dim)  # Input injection
z0 = torch.zeros(batch_size, hidden_dim)  # Initial fixed point estimate

# Create fixed-point function with input injection
f = lambda z: layer(z, x)

# Solve for fixed point
z_out, info = deq(f, z0)
z_star = z_out[-1]  # Final fixed point solution

print(f"Fixed point shape: {z_star.shape}")
print(f"Convergence (abs): {info['abs_lowest'].mean().item():.6f}")
print(f"Solver steps: {info['nstep'].mean().item():.0f}")
```

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

### add_deq_args - CLI Argument Setup

Source: https://context7.com/locuslab/torchdeq/llms.txt

A decorator function that adds commonly used DEQ arguments to an argparse parser.

```APIDOC
## add_deq_args - CLI Argument Setup

### Description
Decorator function that adds all commonly used DEQ arguments to an argparse parser.

### Method
`add_deq_args(parser)`

### Parameters
- **parser**: An argparse.ArgumentParser instance.

### Request Example
```python
import argparse
from torchdeq import add_deq_args, get_deq

parser = argparse.ArgumentParser(description='DEQ Training')
add_deq_args(parser)
args = parser.parse_args()
```
```

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

### Python Argument Parsing for TorchDEQ Configuration

Source: https://context7.com/locuslab/torchdeq/llms.txt

This Python code snippet utilizes the `argparse` module to define and parse command-line arguments for a DEQ model. It includes custom arguments like learning rate and epochs, and integrates automatically generated DEQ arguments. The parsed arguments are then used to initialize the DEQ model, with an example of overriding specific settings.

```python
import argparse

# Assume add_deq_args and get_deq are defined elsewhere
def add_deq_args(parser):
    # Placeholder for adding DEQ specific arguments
    parser.add_argument('--f_solver', type=str, default='broyden')
    parser.add_argument('--b_solver', type=str, default='broyden')
    parser.add_argument('--f_max_iter', type=int, default=100)
    parser.add_argument('--b_max_iter', type=int, default=100)
    parser.add_argument('--f_tol', type=float, default=1e-4)
    parser.add_argument('--b_tol', type=float, default=1e-4)
    parser.add_argument('--f_stop_mode', type=str, default='quick')
    parser.add_argument('--eval_factor', type=float, default=1.0)
    parser.add_argument('--norm_type', type=str, default='none')
    parser.add_argument('--norm_clip', action='store_true')
    parser.add_argument('--norm_clip_value', type=float, default=1.0)
    parser.add_argument('--core', type=str, default='dense')
    parser.add_argument('--ift', action='store_true')
    parser.add_argument('--n_states', type=int, default=1)
    parser.add_argument('--grad', type=str, default='none')
    parser.add_argument('--tau', type=float, default=0.0)
    parser.add_argument('--gamma', type=float, default=1.0)
    parser.add_argument('--jac_loss_weight', type=float, default=0.0)
    parser.add_argument('--jac_loss_freq', type=int, default=1)

def get_deq(args):
    # Placeholder for DEQ model creation
    print(f"Creating DEQ with args: {args}")
    return "DEQ Model"

# Initialize argument parser
parser = argparse.ArgumentParser()

# Add your custom arguments
parser.add_argument('--lr', type=float, default=1e-3)
parser.add_argument('--epochs', type=int, default=100)

# Add all DEQ arguments automatically
add_deq_args(parser)

# Parse arguments (with DEQ defaults)
args = parser.parse_args([])

# Available DEQ arguments after add_deq_args:
print(f"""
DEQ Configuration:
  Solver:
    --f_solver: {args.f_solver}        # Forward solver
    --b_solver: {args.b_solver}        # Backward solver
    --f_max_iter: {args.f_max_iter}    # Forward max iterations
    --b_max_iter: {args.b_max_iter}    # Backward max iterations
    --f_tol: {args.f_tol}              # Forward tolerance
    --b_tol: {args.b_tol}              # Backward tolerance
    --f_stop_mode: {args.f_stop_mode}  # Forward stop mode
    --eval_factor: {args.eval_factor}  # Eval iteration multiplier

  Normalization:
    --norm_type: {args.norm_type}      # Normalization type
    --norm_clip: {args.norm_clip}      # Enable norm clipping
    --norm_clip_value: {args.norm_clip_value}

  Training:
    --core: {args.core}                # DEQ core type
    --ift: {args.ift}                  # Implicit differentiation
    --n_states: {args.n_states}        # Trajectory states
    --grad: {args.grad}                # PhantomGrad steps
    --tau: {args.tau}                  # Damping factor
    --gamma: {args.gamma}              # FP correction gamma

  Regularization:
    --jac_loss_weight: {args.jac_loss_weight}
    --jac_loss_freq: {args.jac_loss_freq}
""")

# Use parsed args to create DEQ
deq = get_deq(args)

# Override specific settings via command line
args2 = parser.parse_args([
    '--f_solver', 'anderson',
    '--f_max_iter', '60',
    '--ift',
    '--n_states', '3',
    '--norm_type', 'spectral_norm'
])
deq2 = get_deq(args2)

```

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

### Jacobian Regularization (jac_reg)

Source: https://context7.com/locuslab/torchdeq/llms.txt

Estimates the Jacobian norm for regularization using Hutchinson estimator to encourage stable fixed points. Includes an example of using the power method to estimate the spectral radius.

```APIDOC
## jac_reg - Jacobian Regularization

### Description
Estimates the Jacobian norm for regularization using Hutchinson estimator to encourage stable fixed points.

### Method
`jac_reg(f_z, z, vecs=2, create_graph=True)`

### Parameters
- **f_z**: Output of the function.
- **z**: Input to the function.
- **vecs**: Number of random vectors for Hutchinson estimator. Defaults to 2.
- **create_graph**: Enable for backprop through the loss. Defaults to True.

### Request Example
```python
with torch.enable_grad():
    f_z = func(z)
    jac_loss = jac_reg(
        f_z,                # Output of function
        z,                  # Input to function
        vecs=2,             # Number of random vectors for Hutchinson estimator
        create_graph=True   # Enable for backprop through the loss
    )
```

### Response
#### Success Response (200)
- **jac_loss** (torch.Tensor) - The calculated Jacobian regularization loss.

### Response Example
```json
{
  "jac_loss": 0.056789
}
```

## power_method - Spectral Radius Estimation

### Description
Estimates the spectral radius of a function's Jacobian using the power method.

### Method
`power_method(f_z, z, n_iters=100)`

### Parameters
- **f_z**: Output of the function.
- **z**: Input to the function.
- **n_iters**: Number of power iterations. Defaults to 100.

### Request Example
```python
eigenvector, eigenvalue = power_method(
    f_z, z,
    n_iters=100  # Number of power iterations
)
```

### Response
#### Success Response (200)
- **eigenvector** (torch.Tensor) - The estimated eigenvector.
- **eigenvalue** (torch.Tensor) - The estimated eigenvalue (spectral radius).

### Response Example
```json
{
  "eigenvector": [0.1, 0.2, ...],
  "eigenvalue": 0.95
}
```
```

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

### Apply Weight Normalization with apply_norm

Source: https://context7.com/locuslab/torchdeq/llms.txt

The apply_norm function automatically applies weight or spectral normalization to DEQ layer weights to enhance training stability. It supports filtering specific layers and includes options for clipping. The example shows applying different normalization types and how to reset or remove normalization.

```python
import torch.nn as nn
from torchdeq.norm import apply_norm, reset_norm, remove_norm

# Define DEQ function layers
class DEQFunc(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_model * 4)
        self.linear2 = nn.Linear(d_model * 4, d_model)
        self.embedding = nn.Embedding(1000, d_model)  # Skip normalization
        self.activation = nn.GELU()

    def forward(self, z, x):
        h = self.activation(self.linear1(z))
        return self.linear2(h) + x

model = DEQFunc(d_model=256)

# Apply weight normalization, skipping embedding layers
apply_norm(
    model,
    norm_type='weight_norm',       # 'weight_norm', 'spectral_norm', or 'none'
    filter_out=['embedding'],       # Skip layers with 'embedding' in name
    norm_clip=True,                 # Clip normalization factor
    norm_clip_value=1.0             # Upper bound for normalization
)

# Apply spectral normalization instead
model2 = DEQFunc(d_model=256)
apply_norm(
    model2,
    norm_type='spectral_norm',
    sn_n_power_iters=1              # Power iterations for spectral norm
)

# Reset normalization before each training iteration
reset_norm(model)

# Remove normalization (e.g., for inference optimization)
remove_norm(model)
```

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

### Initialize DEQ Components

Source: https://github.com/locuslab/torchdeq/blob/main/README.md

Utility functions for adding DEQ-specific arguments to a parser and instantiating a DEQ layer.

```python
add_deq_args(parser)
deq = get_deq(args)
```

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

### Configuration and Initialization

Source: https://context7.com/locuslab/torchdeq/llms.txt

This section describes how to inject DEQ-specific arguments into an argparse parser and initialize the model.

```APIDOC
## Configuration and Initialization

### Description
Integrates DEQ-specific hyperparameters into a standard argparse object and initializes the DEQ model instance.

### Method
N/A (Python Library Function)

### Endpoint
add_deq_args(parser) / get_deq(args)

### Parameters
#### Arguments
- **--f_solver** (string) - Optional - Forward solver algorithm (e.g., anderson)
- **--f_max_iter** (int) - Optional - Maximum iterations for forward solver
- **--ift** (bool) - Optional - Enable implicit differentiation
- **--n_states** (int) - Optional - Number of trajectory states
- **--norm_type** (string) - Optional - Normalization method (e.g., spectral_norm)

### Request Example
parser = argparse.ArgumentParser()
add_deq_args(parser)
args = parser.parse_args(['--f_solver', 'anderson', '--f_max_iter', '60'])
deq = get_deq(args)

### Response
#### Success Response (Object)
- **deq** (object) - Initialized DEQ model instance configured with provided arguments
```

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

### Configure and Initialize DEQ Modules

Source: https://github.com/locuslab/torchdeq/blob/main/docs/get_started.md

Utility functions for adding DEQ-specific arguments to parsers, initializing the DEQ module, and applying normalization to layers.

```python
add_deq_args(parser)
self.deq = get_deq(args)
apply_norm(self.deq_layers)
```

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

### Dataset Directory Structure

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-flow/README.md

Illustrates the expected directory structure for organizing downloaded datasets. Proper organization is crucial for the training and evaluation scripts to locate the data correctly.

```Shell
├── data
    ├── Sintel
        ├── test
        ├── training
    ├── KITTI
        ├── testing
        ├── training
        ├── devkit
    ├── FlyingChairs_release
        ├── data
    ├── FlyingThings3D
        ├── frames_cleanpass
        ├── frames_finalpass
        ├── optical_flow
```

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

### Execute DEQ Forward Pass

Source: https://github.com/locuslab/torchdeq/blob/main/docs/get_started.md

Demonstrates how to define a fixed-point function and execute the DEQ forward pass to obtain equilibrium outputs and convergence information.

```python
def fn(a, b, c):
    return a, b, c
z_out, info = self.deq(fn, (a0, b0, c0))
```

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

### Prepare Wikitext-103 Dataset

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-seq/README.md

Downloads the required Wikitext-103 dataset for training and evaluation using a shell script.

```bash
bash get_data.sh
```

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

### Train and Evaluate IGNN Models

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/ignn/README.md

Commands to initiate training on the PPI dataset and evaluate pretrained models. The evaluation command requires a path to a valid checkpoint file.

```bash
bash run_ignn.sh
bash run_ignn.sh --eval --resume_path [YOUR_CHECKPOINT].pth
```

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

### anderson_solver - Anderson Acceleration

Source: https://context7.com/locuslab/torchdeq/llms.txt

Anderson acceleration solver for faster fixed-point convergence by storing previous residuals.

```APIDOC
## anderson_solver - Anderson Acceleration

### Description
Anderson acceleration solver for faster fixed-point convergence. Stores previous residuals to accelerate the iteration.

### Method
`anderson_solver(func, x0, max_iter=50, tol=1e-6, stop_mode='abs', m=6, lam=1e-4, tau=1.0, return_final=False)`

### Endpoint
N/A (Function call within Python code)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
```python
import torch
from torchdeq.solver import get_solver, anderson_solver

# Get solver by name
solver = get_solver('anderson')

# Define fixed-point function (find sqrt(2) via z = 0.5*(z + 2/z))
f = lambda z: 0.5 * (z + 2.0 / z)
z0 = torch.tensor([[1.0]])  # Initial guess (needs batch dimension)

# Solve with Anderson acceleration
z_star, trajectory, info = anderson_solver(
    f, z0,
    max_iter=50,       # Maximum iterations
    tol=1e-6,          # Convergence tolerance
    stop_mode='abs',   # 'abs' or 'rel' convergence criterion
    m=6,               # Number of stored residuals (memory)
    lam=1e-4,          # Regularization parameter
    tau=1.0,           # Damping factor
    return_final=False # Return best estimate, not final
)

print(f"Solution: {z_star.item():.10f}")
print(f"True sqrt(2): {2**0.5:.10f}")
print(f"Residual: {info['abs_lowest'].item():.2e}")
print(f"Steps: {info['nstep'].item():.0f}")
```

### Response
- **z_star** (torch.Tensor) - The estimated fixed-point solution.
- **trajectory** (list) - A list of intermediate solutions if `return_final=False`.
- **info** (dict) - Dictionary containing convergence information (e.g., 'abs_lowest', 'nstep').

#### Success Response (200)
- **z_star** (torch.Tensor)
- **trajectory** (list)
- **info** (dict)

#### Response Example
```json
{
  "z_star": "tensor([[1.4142]])",
  "trajectory": [],
  "info": {
    "abs_lowest": "0.00e+00",
    "nstep": "10.0"
  }
}
```
```

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

### Download Model Checkpoints

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/ignn/README.md

Uses wget to download specific model checkpoint files from Google Drive using their unique file identifiers.

```bash
wget https://drive.google.com/uc?id=FILE_ID
```

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

### Anderson Solver for Fixed-Point Convergence

Source: https://context7.com/locuslab/torchdeq/llms.txt

Implements the Anderson acceleration solver, which speeds up fixed-point convergence by storing and utilizing previous residuals. This solver is configurable with parameters like maximum iterations, tolerance, convergence stop mode, memory size (m), regularization (lam), and damping (tau).

```python
import torch
from torchdeq.solver import get_solver, anderson_solver

# Get solver by name
solver = get_solver('anderson')

# Define fixed-point function (find sqrt(2) via z = 0.5*(z + 2/z))
f = lambda z: 0.5 * (z + 2.0 / z)
z0 = torch.tensor([[1.0]])  # Initial guess (needs batch dimension)

# Solve with Anderson acceleration
z_star, trajectory, info = anderson_solver(
    f, z0,
    max_iter=50,       # Maximum iterations
    tol=1e-6,          # Convergence tolerance
    stop_mode='abs',   # 'abs' or 'rel' convergence criterion
    m=6,               # Number of stored residuals (memory)
    lam=1e-4,          # Regularization parameter
    tau=1.0,           # Damping factor
    return_final=False # Return best estimate, not final
)

print(f"Solution: {z_star.item():.10f}")
print(f"True sqrt(2): {2**0.5:.10f}")
print(f"Residual: {info['abs_lowest'].item():.2e}")
print(f"Steps: {info['nstep'].item():.0f}")
```

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

### Train MDEQ on CIFAR10

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/mdeq/README.md

Executes the training script for MDEQ on the CIFAR10 dataset. This command should be run from the 'mdeq_cifar' directory.

```bash
bash run.sh
```

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

### get_deq - Create DEQ Module

Source: https://context7.com/locuslab/torchdeq/llms.txt

Factory function to initialize a Deep Equilibrium Model module with configurable solvers and computational graph types.

```APIDOC
## get_deq

### Description
Initializes a DEQ module. It supports different computational graphs (indexing, sliced) and various fixed-point solvers for forward and backward passes.

### Method
Python Factory Function

### Parameters
#### Arguments
- **core** (string) - Optional - Computational graph type: 'indexing' or 'sliced'.
- **f_solver** (string) - Optional - Forward solver: 'anderson', 'broyden', 'fixed_point_iter'.
- **b_solver** (string) - Optional - Backward solver: 'fixed_point_iter'.
- **f_max_iter** (int) - Optional - Maximum iterations for forward pass.
- **ift** (bool) - Optional - Whether to enable Implicit Function Theorem.

### Request Example
```python
deq = get_deq(core='indexing', f_solver='anderson', ift=True)
```

### Response
#### Success Response
- **deq** (nn.Module) - Configured DEQ module instance.
```

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

### Execute DEQ Forward Pass

Source: https://github.com/locuslab/torchdeq/blob/main/README.md

Demonstrates how to perform a forward pass through a DEQ layer using a fixed-point function definition.

```python
def fn(a, b, c):
    return a, b, c

z_out, info = deq(fn, (a0, b0, c0))
```

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

### Fixed Point Iteration Solvers

Source: https://context7.com/locuslab/torchdeq/llms.txt

Provides two fixed-point iteration solvers: `fixed_point_iter` for detailed convergence monitoring and optional damping, and `simple_fixed_point_iter` for a faster, baseline implementation without monitoring. Both solvers aim to find a point `z` where `f(z) = z`.

```python
import torch
from torchdeq.solver import fixed_point_iter, simple_fixed_point_iter

# Fixed-point iteration to find cos(z) = z
f = lambda z: torch.cos(z)
z0 = torch.tensor([[0.5]])

z_star, trajectory, info = fixed_point_iter(
    f, z0,
    max_iter=100,
    tol=1e-8,
    stop_mode='abs',
    tau=1.0,           # Damping: z_new = tau*f(z) + (1-tau)*z
    indexing=[10, 20, 30],  # Store solutions at these iterations
    return_final=False
)

print(f"Fixed point: {z_star.item():.10f}")
print(f"Verification f(z*): {torch.cos(z_star).item():.10f}")
print(f"Trajectory length: {len(trajectory)}")

# Simple (faster) version without convergence monitoring
z_star_simple, _, info_simple = simple_fixed_point_iter(
    f, z0,
    max_iter=100,
    tau=1.0
)
print(f"Simple solver result: {z_star_simple.item():.10f}")
```

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

### Train DEQ-Flow-B Model

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-flow/README.md

Executes the training script for the DEQ-Flow-B model. This process requires two GPUs and follows the standard training pipeline outlined in the project.

```bash
bash train_B.sh
```

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

### TorchDEQ Core Utilities

Source: https://github.com/locuslab/torchdeq/blob/main/docs/get_started.md

Core functions for initializing and configuring Deep Equilibrium models within a PyTorch module.

```APIDOC
## TorchDEQ Core Functions

### Description
These functions are used to initialize the DEQ module, apply automatic normalization to layers, and inject DEQ-specific arguments into an argparse parser.

### Methods
- `add_deq_args(parser)`: Adds DEQ-related command line arguments to an existing parser.
- `get_deq(args)`: Initializes and returns the DEQ module based on provided arguments.
- `apply_norm(layers)`: Automatically applies normalization to the specified DEQ layers.

### Usage Example
```python
from torchdeq import get_deq, apply_norm
from torchdeq.utils import add_deq_args

# Configure arguments
parser = argparse.ArgumentParser()
add_deq_args(parser)
args = parser.parse_args()

# Initialize DEQ
self.deq = get_deq(args)
apply_norm(self.deq_layers)
```
```

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

### Fixed-Point Correction (fp_correction)

Source: https://context7.com/locuslab/torchdeq/llms.txt

Demonstrates how to use the fp_correction function for fixed-point correction with different weighting schemes (exponential, linear, custom) and options for returning individual loss values.

```APIDOC
## fp_correction - Fixed-Point Correction

### Description
Applies fixed-point correction with exponential weighting to a criterion.

### Method
`fp_correction(criterion, (trajectory, target), weight_func='exp', gamma=0.8)`

### Parameters
- **criterion**: The loss function to apply.
- **(trajectory, target)**: A tuple containing predictions and targets.
- **weight_func**: Weighting function ('exp', 'linear', or 'const'). Defaults to 'exp'.
- **gamma**: Decay factor for 'exp' weighting. Defaults to 0.8.

### Request Example
```python
loss = fp_correction(
    criterion,
    (trajectory, target),  # (predictions, targets)
    weight_func='exp',     # 'exp', 'linear', or 'const'
    gamma=0.8              # Decay factor for 'exp'
)
```

### Response
#### Success Response (200)
- **loss** (torch.Tensor) - The calculated loss.
- **loss_values** (list) - Optional: Individual loss values if `return_loss_values=True`.

### Response Example
```json
{
  "loss": 0.1234,
  "loss_values": [0.1, 0.15, 0.12]
}
```

## Custom Weight Function Registration

### Description
Allows registration of custom weight functions for `fp_correction`.

### Method
`register_weight_func(name, func)`

### Parameters
- **name** (str): The name to register the function under.
- **func**: The custom weight function.

### Request Example
```python
def custom_weight(n, k, **kwargs):
    """Quadratic decay weight function."""
    return ((n - k) / n) ** 2

register_weight_func('quadratic', custom_weight)
```
```

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

### Train MDEQ on ImageNet

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/mdeq/README.md

Executes the training script for MDEQ on the ImageNet dataset using PyTorch Distributed Data Parallel (DDP). Requires specifying the data path and optionally the number of GPUs and batch size.

```bash
bash run_mdeq.sh --data YOUR_DATA_PATH
```

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

### Evaluate Pre-trained DEQ Models

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-seq/README.md

Evaluates a pre-trained model checkpoint by providing the path to the file and setting the memory length for inference.

```bash
bash deq_seq_ift.sh --eval --name [NAME_YOUR_RUN] --load_path [CHECKPOINT_NAME].pth
```

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

### fixed_point_iter - Fixed Point Iteration

Source: https://context7.com/locuslab/torchdeq/llms.txt

Simple fixed-point iteration solver with optional damping and convergence monitoring.

```APIDOC
## fixed_point_iter - Fixed Point Iteration

### Description
Simple fixed-point iteration solver with optional damping. Good baseline solver with convergence monitoring.

### Method
`fixed_point_iter(func, x0, max_iter=100, tol=1e-8, stop_mode='abs', tau=1.0, indexing=None, return_final=False)`
`simple_fixed_point_iter(func, x0, max_iter=100, tau=1.0)`

### Endpoint
N/A (Function call within Python code)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
```python
import torch
from torchdeq.solver import fixed_point_iter, simple_fixed_point_iter

# Fixed-point iteration to find cos(z) = z
f = lambda z: torch.cos(z)
z0 = torch.tensor([[0.5]])

z_star, trajectory, info = fixed_point_iter(
    f, z0,
    max_iter=100,
    tol=1e-8,
    stop_mode='abs',
    tau=1.0,           # Damping: z_new = tau*f(z) + (1-tau)*z
    indexing=[10, 20, 30],  # Store solutions at these iterations
    return_final=False
)

print(f"Fixed point: {z_star.item():.10f}")
print(f"Verification f(z*): {torch.cos(z_star).item():.10f}")
print(f"Trajectory length: {len(trajectory)}")

# Simple (faster) version without convergence monitoring
z_star_simple, _, info_simple = simple_fixed_point_iter(
    f, z0,
    max_iter=100,
    tau=1.0
)
print(f"Simple solver result: {z_star_simple.item():.10f}")
```

### Response
- **z_star** (torch.Tensor) - The estimated fixed-point solution.
- **trajectory** (list) - A list of intermediate solutions at specified `indexing` iterations.
- **info** (dict) - Dictionary containing convergence information.

#### Success Response (200)
- **z_star** (torch.Tensor)
- **trajectory** (list)
- **info** (dict)

#### Response Example
```json
{
  "z_star": "tensor([[0.739085]])",
  "trajectory": [
    "tensor([[0.87758]])",
    "tensor([[0.63902]])",
    "tensor([[0.80285]])"
  ],
  "info": {
    "abs_lowest": "1.00e-08",
    "nstep": "30.0"
  }
}
```
```

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

### Train DEQ Models with Implicit Differentiation

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-seq/README.md

Executes training scripts for DEQ models using implicit differentiation (IFT) and optional Jacobian regularization.

```bash
bash deq_seq_ift.sh --name ift
bash deq_seq_ift_jr.sh --name ift_jr
```

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

### Evaluate MDEQ on CIFAR10

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/mdeq/README.md

Evaluates a pre-trained MDEQ checkpoint on the CIFAR10 dataset. Requires specifying the evaluation flag and the path to the checkpoint.

```bash
bash run.sh --eval --test_model PATH_TO_YOUR_CKPT
```

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

### Train DEQ-Flow-H Model

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-flow/README.md

Executes the training script for the DEQ-Flow-H model. This model variant requires three GPUs and is trained using a specific schedule.

```bash
bash train_H.sh
```

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

### Evaluate MDEQ on ImageNet

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/mdeq/README.md

Evaluates a pre-trained MDEQ checkpoint on the ImageNet dataset using PyTorch DDP. Requires specifying the evaluation flag, data path, and checkpoint path. GPU count and batch size can be adjusted.

```bash
bash run_mdeq.sh --eval --data YOUR_DATA_PATH --resume YOUR_CKPT_PATH
```

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

### register_solver - Custom Solver Registration

Source: https://context7.com/locuslab/torchdeq/llms.txt

Allows registration of custom fixed-point solvers for use with DEQ models.

```APIDOC
## register_solver - Custom Solver Registration

### Description
Register custom fixed-point solvers for use with DEQ models.

### Method
`register_solver(name: str, solver_func: callable)`

### Endpoint
N/A (Function call within Python code)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
```python
import torch
from torchdeq.solver import register_solver, get_solver
from torchdeq.solver.utils import solver_stat_from_final_step

def newton_solver(func, x0, max_iter=20, tol=1e-6, **kwargs):
    """Simple Newton-like solver for demonstration."""
    x = x0
    for i in range(max_iter):
        with torch.enable_grad():
            x_input = x.clone().requires_grad_(True)
            fx = func(x_input)
            gx = fx - x_input

            # Approximate Newton step
            x = fx.detach() # This is a simplification, a true Newton step would involve Jacobians

            if gx.abs().max() < tol:
                break

    info = solver_stat_from_final_step(x0, x, i + 1)
    return x, [], info

# Register the custom solver
register_solver('newton', newton_solver)

# Use the custom solver
solver = get_solver('newton')
f = lambda z: 0.5 * (z + 2.0 / z) # Function to find sqrt(2)
z_star, _, info = solver(f, torch.tensor([[1.5]]))
print(f"Newton solver result: {z_star.item():.6f}")
```

### Response
N/A (Registers a solver function globally)

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

#### Response Example
N/A
```

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

### Evaluate Pretrained Model

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-flow/README.md

Evaluates a pretrained DEQ-Flow model. The `--eval` flag enables evaluation mode. A checkpoint can be loaded using `--restore_ckpt CKPT_PATH`, and evaluation datasets are specified with the `--validation` flag (e.g., `kitti`, `sintel`).

```bash
bash val.sh --eval --restore_ckpt CKPT_PATH --validation sintel kitti
```

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

### Invert Images using DEQ-DDIM Script

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/deq-ddim/README.md

Executes a bash script to perform image inversion, transforming images back into the latent space of diffusion models using the DEQ-DDIM approach. This is a key functionality for understanding and manipulating generative models.

```bash
bash scripts/invert_models_deq.sh
```

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

### Optimize ImageNet Evaluation Inference

Source: https://github.com/locuslab/torchdeq/blob/main/deq-zoo/mdeq/README.md

Evaluates MDEQ on ImageNet while optimizing inference speed by limiting the maximum number of solver steps (NFEs). Demonstrates trade-offs between speed and accuracy.

```bash
bash run_mdeq.sh --eval --data YOUR_DATA_PATH --resume YOUR_CKPT_PATH --eval_f_max_iter 20
```

```bash
bash run_mdeq.sh --eval --data YOUR_DATA_PATH --resume YOUR_CKPT_PATH --eval_f_max_iter 10
```

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

### DEQ Forward Pass

Source: https://github.com/locuslab/torchdeq/blob/main/docs/get_started.md

Executing the forward pass for a fixed-point system using the DEQ module.

```APIDOC
## DEQ Forward Pass

### Description
Executes the forward pass of a fixed-point system defined by a callable function. The DEQ module automatically handles the equilibrium finding process.

### Parameters
- **fn** (callable) - Required - A function defining the fixed-point system, taking input tensors and returning tensors of the same shape.
- **z0** (tuple/tensor) - Required - Initial state or input injection for the equilibrium solver.

### Response
- **z_out** (tensor) - The equilibrium output.
- **info** (dict) - Metadata containing convergence information such as 'rel_lowest' and 'abs_lowest'.

### Example
```python
# Define the fixed point system
def fn(a, b, c):
    return a, b, c

# Execute forward pass
z_out, info = self.deq(fn, (a0, b0, c0))
```
```

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

### DEQ Transformer Language Model Implementation

Source: https://context7.com/locuslab/torchdeq/llms.txt

Defines the DEQTransformer model, including embedding, positional encoding, DEQ blocks, and output projection. It handles both training with loss calculation (including fixed-point correction and Jacobian regularization) and inference.

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchdeq import get_deq, reset_deq, add_deq_args
from torchdeq.norm import apply_norm
from torchdeq.dropout import VariationalDropout1d
from torchdeq.loss import fp_correction, jac_reg
import argparse

class TransformerDEQBlock(nn.Module):
    """Single transformer block for DEQ."""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            VariationalDropout1d(dropout),
            nn.Linear(d_ff, d_model),
            VariationalDropout1d(dropout)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, z, x, mask=None):
        # z: current estimate, x: input injection
        z_norm = self.norm1(z)
        attn_out, _ = self.attn(z_norm, z_norm, z_norm, attn_mask=mask)
        z = z + attn_out
        z = z + self.ff(self.norm2(z))
        return z + x  # Input injection

class DEQTransformer(nn.Module):
    """DEQ-based Transformer Language Model."""
    def __init__(self, args, vocab_size, d_model=256, n_heads=4, d_ff=1024):
        super().__init__()
        self.args = args
        self.d_model = d_model

        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = nn.Parameter(torch.randn(1, 512, d_model) * 0.02)

        self.input_proj = nn.Linear(d_model, d_model)
        self.deq_block = TransformerDEQBlock(d_model, n_heads, d_ff)
        self.output_proj = nn.Linear(d_model, vocab_size)

        # Apply weight normalization to DEQ block
        apply_norm(self.deq_block, args=args, filter_out=['norm'])
        self.deq = get_deq(args)

    def forward(self, input_ids, targets=None):
        batch_size, seq_len = input_ids.shape

        # Reset normalization and dropout masks
        reset_deq(self)

        # Embeddings with positional encoding
        x = self.embedding(input_ids)
        x = x + self.pos_encoding[:, :seq_len, :]
        x = self.input_proj(x)  # Input injection

        # Causal mask
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
        mask = mask.to(input_ids.device)

        # Initial state
        z0 = torch.zeros(batch_size, seq_len, self.d_model, device=input_ids.device)

        # DEQ forward: find fixed point z* = f(z*, x)
        f = lambda z: self.deq_block(z, x, mask)
        z_out, info = self.deq(f, z0)

        # Compute loss on trajectory states
        if self.training and targets is not None:
            # Project all trajectory states
            logits_list = [self.output_proj(z) for z in z_out]

            # Fixed-point correction loss
            def cross_entropy(logits, tgt):
                return F.cross_entropy(
                    logits.view(-1, logits.size(-1)),
                    tgt.view(-1)
                )

            loss = fp_correction(
                cross_entropy,
                (logits_list, targets),
                gamma=self.args.gamma
            )

            # Jacobian regularization
            jac_loss = torch.tensor(0.0)
            if self.args.jac_loss_weight > 0:
                z_star = z_out[-1]
                jac_loss = jac_reg(f(z_star), z_star, vecs=1)
                loss = loss + self.args.jac_loss_weight * jac_loss

            return loss, info
        else:
            logits = self.output_proj(z_out[-1])
            return logits, info
```