### Install DASH-pytorch

Source: https://github.com/lucidrains/dash/blob/main/README.md

Install the DASH-pytorch library using pip.

```bash
pip install DASH-pytorch
```

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

### Basic Usage of DASH AdamW Optimizer

Source: https://github.com/lucidrains/dash/blob/main/README.md

Demonstrates initializing and using the DASH AdamW optimizer with a PyTorch model. Ensure parameters and gradients are managed correctly.

```python
import torch
from torch.nn import Linear

from DASH import AdamW

net = Linear(10, 5)
optim = AdamW(net.parameters(), lr = 3e-4)

loss = net(torch.randn(10)).sum()
loss.backward()

optim.step()
optim.zero_grad()

optim.shrink_params()
optim.clear_grad_ema()
```

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

### Initialize and Use DASH AdamW

Source: https://context7.com/lucidrains/dash/llms.txt

Configure the DASH AdamW optimizer and integrate it into a standard training loop, including parameter shrinking and EMA clearing for warm-starts.

```python
import torch
from torch.nn import Linear, Sequential, ReLU

from DASH import AdamW

# Create a simple neural network
net = Sequential(
    Linear(10, 64),
    ReLU(),
    Linear(64, 32),
    ReLU(),
    Linear(32, 5)
)

# Initialize the DASH AdamW optimizer
optim = AdamW(
    net.parameters(),
    lr=3e-4,                  # Learning rate
    betas=(0.9, 0.99),        # Adam beta parameters
    min_shrink_thres=0.3,     # Minimum shrinkage threshold based on cosine similarity
    weight_decay=0.01,        # Decoupled weight decay
    eps=1e-8,                 # Epsilon for numerical stability
    regen_reg_rate=0.,        # Regenerative regularization rate (mutually exclusive with weight_decay)
    grad_ema=True,            # Enable gradient EMA tracking for shrinking
    grad_ema_decay=0.98       # Gradient EMA decay rate
)

# Standard training loop
for epoch in range(10):
    for batch in range(100):
        inputs = torch.randn(32, 10)
        targets = torch.randn(32, 5)

        outputs = net(inputs)
        loss = ((outputs - targets) ** 2).mean()
        loss.backward()

        optim.step()
        optim.zero_grad()

# After training on initial dataset, apply DASH shrinking before warm-starting
# This shrinks parameters based on gradient direction to maintain plasticity
optim.shrink_params()

# Clear gradient EMA before training on new data distribution
optim.clear_grad_ema()

# Continue training on new/updated dataset with maintained plasticity
for epoch in range(10):
    for batch in range(100):
        new_inputs = torch.randn(32, 10)
        new_targets = torch.randn(32, 5)

        outputs = net(new_inputs)
        loss = ((outputs - new_targets) ** 2).mean()
        loss.backward()

        optim.step()
        optim.zero_grad()

print("Training completed with maintained plasticity")
```

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

### Dataset-based Shrinking with DASH

Source: https://context7.com/lucidrains/dash/llms.txt

For dataset-based shrinking, use `shrink_params_with_dataset_()`. This method handles gradient computation and shrinking in a single call, simplifying the process when adapting to new data distributions.

```python
from DASH import AdamW

optimizer = AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01)

# Assuming 'dataloader' provides the dataset for shrinking
optimizer.shrink_params_with_dataset_(dataloader)
```

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

### Replace AdamW with DASH.AdamW

Source: https://context7.com/lucidrains/dash/llms.txt

To integrate DASH, replace `torch.optim.AdamW` with `DASH.AdamW` in your PyTorch workflow. Ensure to call `shrink_params()` and `clear_grad_ema()` at appropriate transition points between training phases.

```python
from DASH import AdamW

# Replace torch.optim.AdamW with DASH.AdamW
optimizer = AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01)

# ... training loop ...

# Call at transition points
optimizer.shrink_params()
optimizer.clear_grad_ema()
```

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

### Control Gradient EMA Tracking with turn_on_grad_ema/turn_off_grad_ema

Source: https://context7.com/lucidrains/dash/llms.txt

Enable or disable gradient EMA tracking at runtime for fine-grained control over gradient statistics collection. Useful for specific training phases or evaluation.

```python
import torch
from torch.nn import Linear
from DASH import AdamW

net = Linear(10, 5)
optim = AdamW(net.parameters(), lr=3e-4, grad_ema=True)

# Initial training with gradient EMA enabled (default)
for _ in range(50):
    loss = net(torch.randn(16, 10)).sum()
    loss.backward()
    optim.step()
    optim.zero_grad()

# Disable gradient EMA during evaluation or specific training phases
optim.turn_off_grad_ema()

for _ in range(20):
    loss = net(torch.randn(16, 10)).sum()
    loss.backward()
    optim.step()  # Steps without updating gradient EMA
    optim.zero_grad()

# Re-enable gradient EMA tracking
optim.turn_on_grad_ema()

# Continue training with gradient EMA collection
for _ in range(50):
    loss = net(torch.randn(16, 10)).sum()
    loss.backward()
    optim.step()
    optim.zero_grad()

# Now shrink_params will use the gradient EMA collected during enabled phases
optim.shrink_params()
```

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

### Shrink Parameters with Dataset

Source: https://github.com/lucidrains/dash/blob/main/README.md

Utilizes the `shrink_params_with_dataset_` function from DASH to adjust model parameters based on a provided dataset.

```python
from DASH import shrink_params_with_dataset_

shrink_params_with_dataset_(net, dataset)
```

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

### Regenerative Regularization

Source: https://context7.com/lucidrains/dash/llms.txt

Configuration for regenerative regularization as an alternative to weight decay.

```APIDOC
## Regenerative Regularization

### Description
Supports regenerative regularization which pulls parameters toward their initial values during training to maintain plasticity. Note that `weight_decay` and `regen_reg_rate` are mutually exclusive.

### Configuration
- **regen_reg_rate** (float) - Required - The rate at which parameters are pulled toward their initial values.
- **weight_decay** (float) - Required - Must be set to 0 when using regenerative regularization.
```

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

### Shrink Parameters with shrink_params_with_dataset_

Source: https://context7.com/lucidrains/dash/llms.txt

Convenience function to shrink parameters using a dataset for gradient statistics. Useful for shrinking based on specific data distributions rather than accumulated training gradients.

```python
import torch
from torch.nn import Linear, Sequential
from torch.utils.data import Dataset
from einops.layers.torch import Reduce

from DASH import AdamW, shrink_params_with_dataset_

# Define a custom dataset
class MyDataset(Dataset):
    def __init__(self, size=1000):
        self.size = size
        self.data = torch.randn(size, 10)

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return self.data[idx]

# Create a network that returns a scalar loss
# (required for automatic backward pass in shrink_params_with_dataset_)
net = Sequential(
    Linear(10, 32),
    torch.nn.ReLU(),
    Linear(32, 1),
    Reduce('b ... -> ', 'sum')  # Reduce to scalar for loss
)

# Train the network normally first
optim = AdamW(net.parameters(), lr=1e-3)
for _ in range(100):
    x = torch.randn(32, 10)
    loss = net(x)
    loss.backward()
    optim.step()
    optim.zero_grad()

# Create dataset for shrinking
shrink_dataset = MyDataset(size=500)

# Shrink parameters using the dataset
# This computes gradient EMA over the dataset and applies direction-aware shrinking
shrink_params_with_dataset_(
    network=net,
    dataset=shrink_dataset,
    optim=optim,           # Optional: uses existing optimizer, creates new one if None
    batch_size=16,         # Batch size for computing gradients
    network_forward_kwargs=dict()  # Additional kwargs passed to network forward
)

print("Parameters shrunk based on dataset gradient statistics")
```

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

### Apply shrink_params Method

Source: https://context7.com/lucidrains/dash/llms.txt

Use shrink_params to adjust parameter magnitudes based on the cosine similarity between negative gradient EMA and parameter values.

```python
import torch
from torch.nn import Linear
from DASH import AdamW

net = Linear(10, 5)
optim = AdamW(net.parameters(), lr=3e-4)

# Train for several iterations to accumulate gradient EMA
for _ in range(50):
    loss = net(torch.randn(16, 10)).sum()
    loss.backward()
    optim.step()
    optim.zero_grad()

# Check parameter magnitudes before shrinking
params_before = [p.clone() for p in net.parameters()]
print(f"Weight norm before shrinking: {net.weight.norm().item():.4f}")

# Apply direction-aware shrinking
# Parameters are multiplied by shrink_factor = max(cosine_similarity(-grad_ema, params), min_shrink_thres)
optim.shrink_params()

print(f"Weight norm after shrinking: {net.weight.norm().item():.4f}")
# Output: Parameters are shrunk based on gradient direction alignment
```

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

### Gradient EMA Control Methods

Source: https://context7.com/lucidrains/dash/llms.txt

Methods to enable or disable gradient EMA tracking at runtime for fine-grained control over gradient statistics collection.

```APIDOC
## turn_on_grad_ema / turn_off_grad_ema

### Description
These methods enable or disable gradient EMA tracking at runtime, allowing fine-grained control over when gradient statistics are collected for the shrinking mechanism.

### Usage
- `optim.turn_off_grad_ema()`: Disables tracking during evaluation or specific training phases.
- `optim.turn_on_grad_ema()`: Re-enables tracking for gradient statistics collection.
```

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

### Apply clear_grad_ema Method

Source: https://context7.com/lucidrains/dash/llms.txt

Reset accumulated gradient EMA statistics to ensure fresh calculations when switching to a new data distribution.

```python
import torch
from torch.nn import Linear
from DASH import AdamW

net = Linear(10, 5)
optim = AdamW(net.parameters(), lr=3e-4, grad_ema=True)

# Train on initial distribution
for _ in range(100):
    loss = net(torch.randn(16, 10)).sum()
    loss.backward()
    optim.step()
    optim.zero_grad()

# Apply shrinking after initial training phase
optim.shrink_params()

# Clear gradient EMA before switching to new data distribution
# This ensures fresh gradient statistics for the new task
optim.clear_grad_ema()
```

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

### Regenerative Regularization with regen_reg_rate

Source: https://context7.com/lucidrains/dash/llms.txt

Utilize regenerative regularization as an alternative to weight decay. This feature pulls parameters toward their initial values during training to maintain plasticity.

```python
import torch
from torch.nn import Linear
from DASH import AdamW

net = Linear(10, 5)

# Use regenerative regularization instead of weight decay
# Note: weight_decay and regen_reg_rate are mutually exclusive
optim = AdamW(
    net.parameters(),
    lr=3e-4,
    weight_decay=0.,       # Must be 0 when using regen_reg_rate
    regen_reg_rate=0.01    # Pulls parameters toward initial values
)

# Initial parameter values are stored automatically
# Training will regularize parameters toward these initial values
for epoch in range(100):
    inputs = torch.randn(32, 10)
    loss = net(inputs).sum()
    loss.backward()

    optim.step()  # Applies regenerative regularization: p.lerp_(param_init, lr/init_lr * regen_rate)
    optim.zero_grad()

print("Training with regenerative regularization completed")
```

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

### shrink_params_with_dataset_

Source: https://context7.com/lucidrains/dash/llms.txt

A convenience function that performs parameter shrinking using a dataset to compute gradient statistics.

```APIDOC
## shrink_params_with_dataset_

### Description
Performs parameter shrinking using a dataset to compute gradient statistics. This is useful when you want to shrink parameters based on gradients computed from a specific dataset rather than accumulated training gradients.

### Parameters
- **network** (torch.nn.Module) - Required - The network to shrink.
- **dataset** (torch.utils.data.Dataset) - Required - The dataset used to compute gradient statistics.
- **optim** (AdamW) - Optional - Existing optimizer instance.
- **batch_size** (int) - Optional - Batch size for computing gradients.
- **network_forward_kwargs** (dict) - Optional - Additional kwargs passed to network forward.
```

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