### Training Example with Gradient Flow in MoEModel

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

A complete training example demonstrating gradient flow through a SinkhornRouter within a Mixture of Experts (MoE) model. This setup enables end-to-end learning.

```python
import torch
from torch import nn
from sinkhorn_router_pytorch import SinkhornRouter

# Model with MoE layer
class MoEModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Linear(256, 512)
        self.router = SinkhornRouter(
            dim=512,
            experts=(8, 512, 512),
            competitive=True,
            causal=False,
            gumbel_noise=True,      # Enable exploration during training
        )
        self.output = nn.Linear(512, 256)

    def forward(self, x):
        x = self.embed(x)
        x = self.router(x)
        return self.output(x)

# Training loop
model = MoEModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.MSELoss()

# Dummy data
x = torch.randn(4, 128, 256)
target = torch.randn(4, 128, 256)

# Forward and backward
model.train()
output = model(x)
loss = criterion(output, target)
loss.backward()

print(f"Loss: {loss.item():.4f}")
print(f"Router expert gradients: {model.router.experts.grad.abs().mean().item():.6f}")

optimizer.step()
```

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

### Install sinkhorn-router-pytorch

Source: https://github.com/lucidrains/sinkhorn-router-pytorch/blob/main/README.md

Install the library using pip. This command is used to add the package to your Python environment.

```bash
pip install sinkhorn-router-pytorch
```

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

### Inference Example with Sinkhorn Router

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

Demonstrates how to use the SinkhornRouter for inference, showing the shapes of routing indices and gate values.

```python
import torch
from sinkhorn_router_pytorch import SinkhornRouter

x = torch.randn(4, 128, 256)
router = SinkhornRouter(dim=256, experts=(8, 256, 256))

# Returns (routing_indices, gate_values) when no experts
routing_indices, gate_values = router(x)

print(f"Routing indices shape: {routing_indices.shape}")
print(f"Gate values shape: {gate_values.shape}")
print(f"Sample routing indices: {routing_indices[0, 0, :10]}")
```

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

### Initialize and Use Sinkhorn Router

Source: https://github.com/lucidrains/sinkhorn-router-pytorch/blob/main/README.md

Demonstrates initializing the SinkhornRouter with expert parameters and then passing input data through it. Ensure the 'experts' parameter is a torch.nn.Parameter.

```python
import torch
from torch import nn
from sinkhorn_router_pytorch import SinkhornRouter

experts = nn.Parameter(torch.randn(8, 8, 512, 256)) # (experts, heads, dim [in], dim [out])

router = SinkhornRouter(
    dim = 512,
    experts = experts,
    competitive = True,
    causal = False,
)

x = torch.randn(1, 8, 1017, 512)
out = router(x) # (1, 8, 1017, 256)
```

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

### Initialize SinkhornRouter with Tuple Expert Initialization

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

Convenient for prototyping, this initializes learnable parameter tensors automatically based on the provided shape tuple. Supports both single-headed and multi-headed experts.

```python
import torch
from sinkhorn_router_pytorch import SinkhornRouter

# Initialize experts with shape tuple: (num_experts, dim_in, dim_out)
# For single-headed experts
router_single = SinkhornRouter(
    dim=512,
    experts=(4, 512, 256),      # Creates nn.Parameter with shape (4, 512, 256)
    competitive=True,
    causal=False,
)

# Initialize multi-headed experts: (num_experts, heads, dim_in, dim_out)
router_multi = SinkhornRouter(
    dim=512,
    experts=(16, 8, 512, 256),  # Creates nn.Parameter with shape (16, 8, 512, 256)
    competitive=True,
    causal=False,
)

# Single-headed forward
x_single = torch.randn(2, 1024, 512)
out_single = router_single(x_single)
print(f"Single-headed output: {out_single.shape}")  # (2, 1024, 256)

# Multi-headed forward
x_multi = torch.randn(2, 8, 1024, 512)
out_multi = router_multi(x_multi)
print(f"Multi-headed output: {out_multi.shape}")  # (2, 8, 1024, 256)
```

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

### Initialize SinkhornRouter with Module-based Experts

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

Use this for traditional MoE feedforward networks. It offers more flexibility when experts have different architectures. Supports variable length masking.

```python
import torch
from torch import nn
from sinkhorn_router_pytorch import SinkhornRouter

# Define experts as a list of modules
experts = [
    nn.Linear(512, 512),
    nn.Linear(512, 512),
    nn.Linear(512, 512),
    nn.Linear(512, 512),
]

# Create router with module-based experts
router = SinkhornRouter(
    dim=512,
    experts=experts,
    competitive=False,          # Non-competitive (uses sigmoid gating)
    causal=False,
)

# Single-headed input: (batch, sequence_length, dim)
x = torch.randn(2, 1024, 512)

# Optional: variable length masking
mask = torch.ones(2, 1024).bool()
mask[0, 800:] = False  # First sequence is shorter
mask[1, 900:] = False  # Second sequence is shorter

out = router(x, mask=mask)
print(f"Output shape: {out.shape}")  # (2, 1024, 512)
```

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

### Initialize SinkhornRouter with Tensor Experts

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

Use this for efficient multi-headed operations. The router handles sequence padding and variable length masking. It can operate in competitive or non-competitive modes.

```python
import torch
from torch import nn
from sinkhorn_router_pytorch import SinkhornRouter

# Define experts as a tensor: (num_experts, heads, dim_in, dim_out)
experts = nn.Parameter(torch.randn(8, 8, 512, 256))

# Create the router with competitive gating (non-causal)
router = SinkhornRouter(
    dim=512,                    # Input dimension
    experts=experts,            # Expert weights tensor or ModuleList
    competitive=True,           # Use competitive gating (default for non-causal)
    causal=False,               # Non-causal routing
    sinkhorn_iters=8,           # Number of Sinkhorn iterations
    temperature=1.0,            # Softmax temperature
    gumbel_noise=False,         # Add Gumbel noise for exploration
)

# Input: (batch, heads, sequence_length, dim)
x = torch.randn(1, 8, 1017, 512)

# Forward pass returns routed output
out = router(x)  # Output shape: (1, 8, 1017, 256)
print(f"Input shape: {x.shape}")
print(f"Output shape: {out.shape}")
```

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

### Sinkhorn Balancing Algorithm

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

Applies the Sinkhorn-Knopp algorithm for balanced routing probabilities. Can optionally add Gumbel noise for stochastic routing during training.

```python
import torch
from sinkhorn_router_pytorch import sinkhorn

# Create unnormalized gate logits: (heads, tokens, experts)
gate_logits = torch.randn(8, 1024, 16)

# Apply Sinkhorn normalization
balanced_gates = sinkhorn(
    gate_logits,
    num_iters=8,                # Number of normalization iterations
    temperature=1.0,            # Temperature for softmax scaling
    gumbel=False,               # Add Gumbel noise for stochastic routing
    noise_inv_temp=1.0,         # Inverse temperature for Gumbel noise
    eps=1e-10,                  # Numerical stability epsilon
)

print(f"Input gates shape: {gate_logits.shape}")
print(f"Balanced gates shape: {balanced_gates.shape}")
print(f"Row sums (should be ~1): {balanced_gates[0].sum(dim=-1)[:5]}")
print(f"Column sums (should be ~1): {balanced_gates[0].sum(dim=-2)[:5]}")

# With Gumbel noise for exploration during training
stochastic_gates = sinkhorn(
    gate_logits,
    num_iters=8,
    gumbel=True,                # Enable Gumbel noise
    noise_inv_temp=0.5,         # Control noise magnitude
)
```

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

### SwitchHead MoE Attention Layer

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

Implements Mixture-of-Experts attention for both value and output projections. Supports causal masking and padding masks for variable sequence lengths.

```python
import torch
from sinkhorn_router_pytorch.switchhead import SwitchHead

# Create SwitchHead attention layer
attn = SwitchHead(
    dim=512,                    # Model dimension
    num_experts=16,             # Number of routing experts
    heads=8,                    # Number of attention heads
    dim_head=64,                # Dimension per head
    dropout=0.1,                # Attention dropout
    causal=True,                # Causal masking for autoregressive
)

# Input: (batch, sequence_length, dim)
x = torch.randn(2, 1017, 512)

# Forward pass through MoE attention
out = attn(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {out.shape}")  # (2, 1017, 512)

# With padding mask for variable length sequences
mask = torch.ones(2, 1017).bool()
mask[0, 800:] = False
out_masked = attn(x, mask=mask)
```

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

### Causal Routing with SinkhornRouter

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

Enables causal masking for autoregressive models. Ensure 'competitive' is set to False when 'causal' is True.

```python
import torch
from torch import nn
from sinkhorn_router_pytorch import SinkhornRouter

experts = nn.Parameter(torch.randn(8, 512, 512))

router = SinkhornRouter(
    dim=512,
    experts=experts,
    causal=True,                # Enable causal masking
    competitive=False,          # Required: causal cannot be competitive
    sinkhorn_iters=8,
    temperature=0.5,            # Lower temperature for sharper routing
)

# Generate sequence autoregressively
x = torch.randn(1, 2048, 512)
out = router(x)
print(f"Causal output shape: {out.shape}")  # (1, 2048, 512)
```

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

### Router-Only Mode without Experts

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

Instantiates SinkhornRouter to return only routing indices and gate values, useful for custom expert implementations or analysis. Specify 'num_experts' without providing 'experts'.

```python
import torch
from sinkhorn_router_pytorch import SinkhornRouter

# Router without experts - returns indices and gate values
router = SinkhornRouter(
    dim=512,
    num_experts=8,              # Specify number of experts
    # experts=None,             # No experts provided
    heads=4,
    competitive=True,
    causal=False,
)

x = torch.randn(2, 4, 1024, 512)

```

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

### Gating Module for External Gate Computation

Source: https://context7.com/lucidrains/sinkhorn-router-pytorch/llms.txt

Utilizes a separate Gating module to compute gate logits, allowing shared gates across multiple routers. Disable internal gating in SinkhornRouter by setting 'has_gating=False'.

```python
import torch
from torch import nn
from sinkhorn_router_pytorch import SinkhornRouter, Gating

# Create separate gating module
gating = Gating(
    dim=512,
    num_experts=16,
    heads=8,
)

# Create routers without internal gating
value_router = SinkhornRouter(
    dim=512,
    experts=nn.Parameter(torch.randn(16, 8, 512, 256)),
    causal=True,
    has_gating=False,           # Disable internal gating
)

output_router = SinkhornRouter(
    dim=512,
    experts=nn.Parameter(torch.randn(16, 8, 256, 512)),
    causal=True,
    has_gating=False,
)

# Compute shared gates once
x = torch.randn(2, 8, 1024, 512)
gates = gating(x)               # Shape: (2, 8, 1024, 16)

# Use same gates for both routers
values = value_router(x, gates=gates)
outputs = output_router(values, gates=gates)

print(f"Shared gates shape: {gates.shape}")
print(f"Values shape: {values.shape}")
print(f"Outputs shape: {outputs.shape}")
```

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