### Basic LFQ Setup

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/lfq.md

Initialize the LFQ module with specified dimensions and loss weights for entropy and commitment.

```python
lfq = LFQ(
    dim=256,
    entropy_loss_weight=0.1,
    commitment_loss_weight=0.25
)
```

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

### Install vector-quantize-pytorch

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Install the library using pip.

```bash
pip install vector-quantize-pytorch
```

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

### Training Vector Quantize Model

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/README.md

Example of a training loop for VectorQuantize, including optimizer setup, forward pass, reconstruction loss calculation, and backpropagation.

```python
from vector_quantize_pytorch import VectorQuantize
import torch.nn.functional as F

vq = VectorQuantize(dim=256, codebook_size=512)
optimizer = torch.optim.Adam(vq.parameters(), lr=1e-3)

# Training loop
for x in dataloader:
    optimizer.zero_grad()
    
    quantized, indices = vq(x)
    
    # Reconstruction loss
    recon_loss = F.mse_loss(quantized, x)
    
    # Get total loss (includes commitment + codebook diversity)
    recon_loss.backward()
    
    optimizer.step()
```

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

### Standard ResidualLFQ Setup

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/residual-lfq.md

Instantiate ResidualLFQ with specified dimensions and loss weights. This is a common setup for general use.

```python
rlfq = ResidualLFQ(
    dim=256,
    num_quantizers=8,
    entropy_loss_weight=0.1,
    commitment_loss_weight=0.25
)
```

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

### FSQ Configuration Examples

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/fsq.md

Illustrates various ways to configure the FSQ class for different quantization strategies and data formats.

```APIDOC
## Configuration Examples

### Simple Symmetric Quantization
Configures FSQ with a specified number of dimensions and levels per dimension, enabling symmetry preservation.
```python
fsq = FSQ(
    levels=[8] * 8,  # 8 dimensions, 8 levels each
    dim=512,
    preserve_symmetry=True
)
```

### Multi-Codebook
Sets up FSQ with multiple independent quantizers, each having its own set of levels.
```python
fsq = FSQ(
    levels=[16, 16, 16],
    dim=256,
    num_codebooks=4  # 4 independent quantizers
)
```

### Different Quantization per Dimension
Defines FSQ with varying numbers of levels for each dimension, allowing for decreasing precision.
```python
fsq = FSQ(
    levels=[32, 16, 16, 8, 8, 8, 8, 8],  # Decreasing precision
    preserve_symmetry=True
)
```

### Channel-First Format
Initializes FSQ for data where the channel dimension comes before the sequence length, common in some deep learning frameworks.
```python
fsq = FSQ(
    levels=[8, 8, 8],
    dim=256,
    channel_first=True  # Input: (batch, 256, seq_len)
)

x = torch.randn(4, 256, 1024)
quantized, indices = fsq(x)
```
```

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

### LFQ Initialization and Forward Pass Example

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/lfq.md

Demonstrates how to initialize the LFQ model and perform a forward pass with sample input. Shows how to access quantized output, indices, and entropy loss. Includes a typical training loss calculation.

```python
from vector_quantize_pytorch import LFQ

lfq = LFQ(
    dim=256,
    entropy_loss_weight=0.1,
    commitment_loss_weight=0.25
)

x = torch.randn(4, 1024, 256)
quantized, indices, entropy_loss = lfq(x)

print(quantized.shape)  # torch.Size([4, 1024, 256])
print(indices.shape)    # torch.Size([4, 1024])
print(entropy_loss)     # scalar tensor

# Typical training loss
loss = entropy_loss + commitment_weight * F.mse_loss(quantized, x)
loss.backward()
```

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

### LFQ Initialization

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/lfq.md

Demonstrates how to initialize the LFQ module with various configurations, including basic setup, different loss weights, multiple codebooks, spherical variant, and soft clamping.

```APIDOC
## LFQ Initialization

### Description
Initializes the LFQ module with specified dimensions and optional parameters for loss weighting, codebook configuration, and quantization behavior.

### Parameters
- **dim** (int) - Required - The input dimension.
- **num_codebooks** (int) - Optional - The number of codebooks to use. Defaults to 1.
- **entropy_loss_weight** (float) - Optional - Weight for the entropy loss. Defaults to 0.1.
- **commitment_loss_weight** (float) - Optional - Weight for the commitment loss. Defaults to 0.25.
- **spherical** (bool) - Optional - Whether to use the spherical variant. Defaults to False.
- **soft_clamp_input_value** (float) - Optional - Value for soft clamping of inputs. If None, soft clamping is disabled. Defaults to None.
- **return_loss_breakdown** (bool) - Optional - Whether to return a detailed loss breakdown. Defaults to False.

### Configuration Examples

#### Basic Setup
```python
LFQ(
    dim=256,
    entropy_loss_weight=0.1,
    commitment_loss_weight=0.25
)
```

#### Low Entropy Loss
```python
LFQ(
    dim=256,
    entropy_loss_weight=0.01,
    commitment_loss_weight=0.25
)
```

#### Multiple Codebooks
```python
LFQ(
    dim=256,
    num_codebooks=4
)
```

#### Spherical Variant
```python
LFQ(
    dim=256,
    spherical=True,
    entropy_loss_weight=0.1
)
```

#### Soft Clamping
```python
LFQ(
    dim=256,
    soft_clamp_input_value=20.0,
    entropy_loss_weight=0.1
)
```
```

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

### HierarchicalVQ Initialization and Forward Pass Example

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/hierarchical-vq.md

Demonstrates how to initialize HierarchicalVQ with specified dimensions, codebook sizes, and number of layers. Shows a typical forward pass with random input and prints the shapes of the output, indices list, and loss.

```python
from vector_quantize_pytorch import HierarchicalVQ

hvq = HierarchicalVQ(
    dim=256,
    codebook_size=512,
    num_layers=3,
    commitment_weight=0.25
)

x = torch.randn(4, 1024, 256)
quantized, indices_list, loss = hvq(x)

print(quantized.shape)  # torch.Size([4, 1024, 256])
print(len(indices_list))  # 3 (one per layer)
print(indices_list[0].shape)  # torch.Size([4, 1024])
```

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

### Spherical Variant LFQ Setup

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/lfq.md

Set up the spherical variant of LFQ, which is based on the paper 'https://arxiv.org/abs/2406.07548', with a specified entropy loss weight.

```python
lfq = LFQ(
    dim=256,
    spherical=True,  # From https://arxiv.org/abs/2406.07548
    entropy_loss_weight=0.1
)
```

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

### Initialize and Use ResidualLFQ

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/residual-lfq.md

Instantiate the ResidualLFQ model and perform a forward pass. This example demonstrates how to set up the model with specified dimensions and weights, and then use it to quantize an input tensor. It also shows how to calculate the training loss using the returned components.

```python
from vector_quantize_pytorch import ResidualLFQ

import torch
import torch.nn.functional as F

rlfq = ResidualLFQ(
    dim=256,
    num_quantizers=8,
    entropy_loss_weight=0.1,
    commitment_loss_weight=0.25
)

x = torch.randn(4, 1024, 256)
quantized, indices, entropy_loss = rlfq(x)

print(quantized.shape)   # torch.Size([4, 1024, 256])
print(indices.shape)     # torch.Size([4, 1024, 8])
print(entropy_loss)      # scalar

# Training loss
commitment_weight = 0.25
loss = entropy_loss + commitment_weight * F.mse_loss(quantized, x)
```

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

### LatentQuantize with Compression

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/latent-quantize.md

Configure LatentQuantize for significant compression by setting a smaller latent_dim. This example uses 3 encoder and decoder layers.

```python
lq = LatentQuantize(
    dim=256,
    codebook_size=512,
    latent_dim=64,  # Compress to 64D for quantization
    encoder_layers=3,
    decoder_layers=3
)
```

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

### ResidualLFQ Spherical Variant Setup

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/residual-lfq.md

Instantiate the spherical variant of ResidualLFQ. This configuration is suitable for applications where spherical quantization properties are beneficial.

```python
rlfq = ResidualLFQ(
    dim=256,
    num_quantizers=8,
    spherical=True,
    entropy_loss_weight=0.1
)
```

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

### Get Quantized Output from Indices

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/vector-quantize.md

Shows how to obtain the quantized output directly from codebook indices. This method applies projection if necessary.

```python
indices = torch.tensor([[0, 1, 2], [3, 4, 5]])
quantized_output = vq.get_output_from_indices(indices)
```

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

### Initialize Grouped Residual Feature Quantizer for Images

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/grouped-quantizers.md

Instantiate GroupedResidualFSQ for image feature processing. This setup uses 2 groups, 4 quantizers per group, and specifies the number of levels for each quantizer (8, 6, 5) for 256-dimensional inputs.

```python
grfsq = GroupedResidualFSQ(
    dim=256,
    groups=2,
    num_quantizers=4,
    levels=[8, 6, 5]
)
```

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

### Initialize and Use LFQ for Image Features

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Demonstrates how to initialize and use the LFQ quantizer with image features. Ensure the codebook size is a power of 2. Experiment with the inverse temperature for optimal results.

```python
import torch
from vector_quantize_pytorch import LFQ

# you can specify either dim or codebook_size
# if both specified, will be validated against each other

quantizer = LFQ(
    codebook_size = 65536,      # codebook size, must be a power of 2
    dim = 16,                   # this is the input feature dimension, defaults to log2(codebook_size) if not defined
    entropy_loss_weight = 0.1,  # how much weight to place on entropy loss
    diversity_gamma = 1.        # within entropy loss, how much weight to give to diversity of codes, taken from https://arxiv.org/abs/1911.05894
)

image_feats = torch.randn(1, 16, 32, 32)

quantized, indices, entropy_aux_loss = quantizer(image_feats, inv_temperature=100.)  # you may want to experiment with temperature

# (1, 16, 32, 32), (1, 32, 32), ()

assert (quantized == quantizer.indices_to_codes(indices)).all()
```

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

### ResidualFSQ num_quantizers property

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/residual-fsq.md

Gets the number of residual quantizers used in the module.

```APIDOC
## num_quantizers

### Description
Number of residual quantizers.

### Property
`num_quantizers`

### Returns
- `count` (int) - The number of quantizers
```

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

### get_output_from_indices Method

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/vector-quantize.md

Gets the quantized output directly from indices, applying projection if necessary.

```APIDOC
## get_output_from_indices

### Description
Get quantized output directly from indices (applies projection if needed).

### Method
`get_output_from_indices(self, indices: Tensor) -> Tensor`

### Parameters:

- **indices** (Tensor): Codebook indices.

### Returns:

- Projected output tensor.

### Example:
```python
indices = torch.tensor([[0, 1, 2], [3, 4, 5]])
output = vq.get_output_from_indices(indices)
print(output.shape) # Example shape: torch.Size([2, 3, 256])
```
```

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

### Access and Set Codebook Embeddings

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/vector-quantize.md

Demonstrates how to get and set the codebook embeddings. The shape of the codebook depends on whether it is multi-headed.

```python
vq = VectorQuantize(dim=256, codebook_size=512)
# Get codebook
codes = vq.codebook  # shape: (512, 256)

# Set codebook
vq.codebook = torch.randn(512, 256)
```

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

### Random Projection Quantizer Initialization

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Initialize a Random Projection Quantizer. This method does not require learning the quantizer and was used in Google's Universal Speech Model.

```python
import torch
from vector_quantize_pytorch import RandomProjectionQuantizer

quantizer = RandomProjectionQuantizer(
    dim = 512,               # input dimensions
    num_codebooks = 16,      # in USM, they used up to 16 for 5% gain
    codebook_dim = 256,      # codebook dimension
    codebook_size = 1024     # codebook size
)

x = torch.randn(1, 1024, 512)
indices = quantizer(x)

# (1, 1024, 16)
```

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

### Initialize and Use LFQ with Multiple Codebooks

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Demonstrates initializing LFQ with multiple codebooks. The total codebook dimension is calculated based on the codebook size and the number of codebooks. The output shape matches the input image features.

```python
import torch
from vector_quantize_pytorch import LFQ

quantizer = LFQ(
    codebook_size = 4096,
    dim = 16,
    num_codebooks = 4  # 4 codebooks, total codebook dimension is log2(4096) * 4
)

image_feats = torch.randn(1, 16, 32, 32)

quantized, indices, entropy_aux_loss = quantizer(image_feats)

# (1, 16, 32, 32), (1, 32, 32, 4), ()

assert image_feats.shape == quantized.shape
assert (quantized == quantizer.indices_to_codes(indices)).all()
```

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

### Get FSQ Codebook Size

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/fsq.md

Demonstrates how to access the total number of discrete codes available in an FSQ instance. This is calculated as the product of all specified levels.

```python
fsq = FSQ(levels=[8, 6, 5])
print(fsq.codebook_size)  # 240
```

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

### Initialize and Use Residual LFQ for Audio Compression

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Shows how to initialize and use ResidualLFQ for audio compression tasks. The output shape of the quantized data matches the input, and the `get_output_from_indices` method can reconstruct the output from indices.

```python
import torch
from vector_quantize_pytorch import ResidualLFQ

residual_lfq = ResidualLFQ(
    dim = 256,
    codebook_size = 256,
    num_quantizers = 8
)

x = torch.randn(1, 1024, 256)

residual_lfq.eval()

quantized, indices, commit_loss = residual_lfq(x)

# (1, 1024, 256), (1, 1024, 8), (8)

quantized_out = residual_lfq.get_output_from_indices(indices)

# (1, 1024, 256)

assert torch.all(quantized == quantized_out)
```

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

### ResidualLFQ with Decreasing Scale per Layer

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/residual-lfq.md

Configure ResidualLFQ to automatically halve the codebook scale at deeper layers. This setup is useful for achieving more compression in later stages.

```python
# Compress more at deeper layers
rlfq = ResidualLFQ(
    dim=256,
    num_quantizers=8,
    codebook_scale=1.0,  # Automatically halved per layer in residual
    entropy_loss_weight=0.1
)
```

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

### Initialize and Use ResidualVQ

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/residual-vq.md

Initializes the ResidualVQ model and demonstrates basic encoding with different return options. Use this to perform the main quantization and reconstruction tasks.

```python
from vector_quantize_pytorch import ResidualVQ

rvq = ResidualVQ(
    dim=256,
    num_quantizers=8,
    codebook_size=1024,
    decay=0.95
)

x = torch.randn(4, 1024, 256)

# Basic encoding
quantized, indices, commit_loss = rvq(x)
print(quantized.shape)   # torch.Size([4, 1024, 256])
print(indices.shape)     # torch.Size([4, 1024, 8])
print(commit_loss.shape) # torch.Size([4, 8])

# Get all intermediate representations
q, idx, loss, all_codes = rvq(x, return_all_codes=True)
print(all_codes.shape)   # torch.Size([8, 4, 1024, 256])
```

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

### Initialize SimVQ

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/sim-vq.md

Instantiate the SimVQ class with specified dimensions and codebook size. Commitment weight can be adjusted for loss balancing. Ensure PyTorch is imported.

```python
from vector_quantize_pytorch import SimVQ

sim_vq = SimVQ(
    dim=256,
    codebook_size=512,
    commitment_weight=0.25
)
```

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

### LFQ Encode Method

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/lfq.md

Shows how to use the `encode` method to get quantization indices from an input tensor without computing gradients. This is useful for inference or when only the discrete codes are needed.

```python
indices = lfq.encode(x)
```

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

### Monitor Loss Components

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/errors.md

When `return_loss_breakdown=True` is enabled, this snippet shows how to monitor individual components of the vector quantization loss, such as commitment, diversity, and orthogonal regularization. Analyzing these can help diagnose numerical issues or guide hyperparameter tuning.

```python
# If return_loss_breakdown=True
quantized, breakdown = vq(x, indices=target, return_loss_breakdown=True)
print(f"Commitment: {breakdown.commitment}")
print(f"Diversity: {breakdown.codebook_diversity}")
print(f"Orthogonal: {breakdown.orthogonal_reg}")
```

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

### Use LFQ for Sequence and Video Features

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Shows how to apply LFQ to sequence data (batch, seq, feat) and video data (batch, feat, time, height, width). The output shape matches the input shape for both data types.

```python
import torch
from vector_quantize_pytorch import LFQ

quantizer = LFQ(
    codebook_size = 65536,
    dim = 16,
    entropy_loss_weight = 0.1,
    diversity_gamma = 1.
)

seq = torch.randn(1, 32, 16)
quantized, *_ = quantizer(seq)

assert seq.shape == quantized.shape

video_feats = torch.randn(1, 16, 10, 32, 32)
quantized, *_ = quantizer(video_feats)

assert video_feats.shape == quantized.shape
```

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

### LFQ Constructor

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/lfq.md

This section details the parameters accepted by the LFQ constructor, allowing users to customize its behavior and configuration.

```APIDOC
## LFQ Constructor

### Description
Initializes the LFQ module with specified configuration parameters.

### Parameters
#### Constructor Parameters
- **dim** (int | None) - Optional - Input feature dimension.
- **codebook_size** (int | None) - Optional - Target codebook size (must be power of 2, inferred from dim if not set).
- **entropy_loss_weight** (float) - Optional - Default: 0.1 - Weight of entropy loss to encourage utilization.
- **commitment_loss_weight** (float) - Optional - Default: 0.0 - MSE loss weight between input and quantized.
- **diversity_gamma** (float) - Optional - Default: 1.0 - Gamma parameter for diversity loss.
- **straight_through_activation** (nn.Module) - Optional - Default: Identity - Activation for straight-through estimator.
- **num_codebooks** (int) - Optional - Default: 1 - Number of independent LFQ codebooks.
- **keep_num_codebooks_dim** (bool | None) - Optional - Keep codebook dimension in output.
- **codebook_scale** (float) - Optional - Default: 1.0 - Scale factor for codebook (useful for residual stacking).
- **frac_per_sample_entropy** (float) - Optional - Default: 1.0 - Fraction of probabilities to use for per-sample entropy (0-1).
- **has_projections** (bool | None) - Optional - Auto-detect if projections needed.
- **projection_has_bias** (bool) - Optional - Default: True - Bias in input/output projections.
- **soft_clamp_input_value** (float | None) - Optional - Soft clamp input to this range.
- **cosine_sim_project_in** (bool) - Optional - Default: False - Use cosine similarity in projection.
- **cosine_sim_project_in_scale** (float | None) - Optional - Scale for cosine similarity projection.
- **channel_first** (bool | None) - Optional - Input format flag.
- **experimental_softplus_entropy_loss** (bool) - Optional - Default: False - Use softplus for entropy loss.
- **entropy_loss_offset** (float) - Optional - Default: 5.0 - Offset before softplus.
- **spherical** (bool) - Optional - Default: False - Spherical quantization variant.
- **force_quantization_f32** (bool) - Optional - Default: True - Force quantization in float32.
- **orthogonal_rotation** (bool) - Optional - Default: False - Apply orthogonal rotation for utilization.
```

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

### LFQ Class Initialization

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/lfq.md

Initializes the LFQ module with various configuration options for quantization and learning.

```APIDOC
## Class LFQ

### Description
Initializes the LFQ module with various configuration options for quantization and learning.

### Parameters

#### `__init__` Parameters

- **dim** (int | None) - The dimension of the input features.
- **codebook_size** (int | None) - The number of codes in the codebook.
- **entropy_loss_weight** (float) - Weight for the entropy loss. Defaults to 0.1.
- **commitment_loss_weight** (float) - Weight for the commitment loss. Defaults to 0.0.
- **diversity_gamma** (float) - Gamma parameter for diversity loss. Defaults to 1.0.
- **straight_through_activation** (nn.Module) - Activation function for straight-through estimation. Defaults to nn.Identity().
- **num_codebooks** (int) - Number of codebooks to use. Defaults to 1.
- **keep_num_codebooks_dim** (bool | None) - Whether to keep the codebook dimension.
- **codebook_scale** (float) - Scale factor for the codebook. Defaults to 1.0.
- **frac_per_sample_entropy** (float) - Fraction of samples to use for entropy calculation. Defaults to 1.0.
- **has_projections** (bool | None) - Whether to use projections.
- **projection_has_bias** (bool) - Whether projections have a bias. Defaults to True.
- **soft_clamp_input_value** (float | None) - Value for soft clamping of input.
- **cosine_sim_project_in** (bool) - Whether to use cosine similarity for input projection. Defaults to False.
- **cosine_sim_project_in_scale** (float | None) - Scale for cosine similarity input projection.
- **channel_first** (bool | None) - Whether the input is channel-first.
- **experimental_softplus_entropy_loss** (bool) - Whether to use experimental softplus for entropy loss. Defaults to False.
- **entropy_loss_offset** (float) - Offset for entropy loss. Defaults to 5.0.
- **spherical** (bool) - Whether to use spherical quantization. Defaults to False.
- **force_quantization_f32** (bool) - Whether to force quantization to f32. Defaults to True.
- **orthogonal_rotation** (bool) - Whether to use orthogonal rotation. Defaults to False.
```

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

### Simple and Fast FSQ Configuration

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/configuration.md

A straightforward and fast configuration for FSQ, using a specified number of levels and preserving symmetry.

```python
FSQ(
    levels=[8] * 8,
    dim=512,
    preserve_symmetry=True
)
```

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

### Sim VQ Initialization with Rotation Trick

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Initialize Sim VQ with the rotation trick. This method uses a frozen codebook and implicitly generated codes via linear projection, aiming to reduce codebook collapse and improve convergence.

```python
import torch
from vector_quantize_pytorch import SimVQ

sim_vq = SimVQ(
    dim = 512,
    codebook_size = 1024,
    rotation_trick = True  # use rotation trick from Fifty et al.
)

x = torch.randn(1, 1024, 512)
quantized, indices, commit_loss = sim_vq(x)

assert x.shape == quantized.shape
assert torch.allclose(quantized, sim_vq.indices_to_codes(indices), atol = 1e-6)
```

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

### HierarchicalVQ Configuration with Varying Codebook Sizes

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/hierarchical-vq.md

Illustrates initializing HierarchicalVQ with a specific number of layers and a tuple defining decreasing codebook sizes for each layer. This allows for finer-grained quantization at deeper levels of the hierarchy.

```python
hvq = HierarchicalVQ(
    dim=256,
    codebook_size=512,
    num_layers=4,
    layer_codebook_size=(1024, 512, 256, 128),  # Decreasing
    commitment_weight=0.25
)
```

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

### ResidualVQ Initialization

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/residual-vq.md

Demonstrates the initialization of the ResidualVQ class with various configuration options such as shared codebooks, quantizer dropout, variable codebook sizes, beam search, and neural codebooks.

```APIDOC
## ResidualVQ Initialization Options

This section details various ways to initialize the `ResidualVQ` class, showcasing different configurations for advanced usage.

### Shared Codebook (RQ-VAE)
```python
rvq = ResidualVQ(
    dim=256,
    num_quantizers=8,
    codebook_size=1024,
    shared_codebook=True,
    stochastic_sample_codes=True,
    sample_codebook_temp=0.1
)
```

### Quantizer Dropout (EnCodec-style)
```python
rvq = ResidualVQ(
    dim=256,
    num_quantizers=8,
    codebook_size=1024,
    quantize_dropout=True,
    quantize_dropout_cutoff_index=1,
    quantize_dropout_multiple_of=4  # Drop 4 at a time
)
```

### Variable Codebook Sizes
```python
rvq = ResidualVQ(
    dim=256,
    codebook_size=(512, 512, 1024, 1024, 2048, 2048, 2048, 2048),
    # Codebook sizes increase for deeper layers
)
print(rvq.num_quantizers)  # 8
```

### Beam Search Decoding
```python
rvq = ResidualVQ(
    dim=256,
    num_quantizers=8,
    codebook_size=1024,
    beam_size=16,        # Training beam search
    eval_beam_size=32    # Inference beam search (higher quality)
)
```

### Neural Codebook (QINCo)
```python
rvq = ResidualVQ(
    dim=256,
    num_quantizers=8,
    codebook_size=1024,
    implicit_neural_codebook=True,
    mlp_kwargs=dict(
        dim_hidden=512,
        depth=4,
        l2norm_output=False
    )
)
```

### Custom VectorQuantize Configuration
```python
rvq = ResidualVQ(
    dim=256,
    num_quantizers=8,
    codebook_size=1024,
    # Pass VectorQuantize kwargs
    commitment_weight=0.25,
    orthogonal_reg_weight=0.1,
    use_cosine_sim=False,
    kmeans_init=True
)
```

### Processing Image Feature Maps
```python
rvq = ResidualVQ(
    dim=256,
    num_quantizers=8,
    codebook_size=1024,
    accept_image_fmap=True
)

# Input: (batch, channels, height, width)
x = torch.randn(4, 256, 32, 32)
quantized, indices, commit_loss = rvq(x)
print(indices.shape)  # torch.Size([4, 32, 32, 8])
```
```

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

### VectorQuantize Constructor

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/vector-quantize.md

Initializes the VectorQuantize module with specified parameters. This module is used for vector quantization, a technique for learning a discrete representation of continuous data.

```APIDOC
## VectorQuantize Constructor

### Description
Initializes the VectorQuantize module with specified parameters. This module is used for vector quantization, a technique for learning a discrete representation of continuous data.

### Parameters
#### Constructor Parameters
- **dim** (int) - Required - Input feature dimension
- **codebook_size** (int) - Required - Number of codebook entries (vocabulary size)
- **codebook_dim** (int | None) - Optional - Dimension of codebook codes (defaults to dim)
- **heads** (int) - Optional - Default: 1 - Number of separate codebook heads
- **separate_codebook_per_head** (bool) - Optional - Default: False - Use independent codebooks per head vs shared codebook
- **decay** (float) - Optional - Default: 0.8 - EMA decay factor for codebook updates (lower = faster change)
- **eps** (float) - Optional - Default: 1e-5 - Epsilon for numerical stability
- **freeze_codebook** (bool) - Optional - Default: False - Prevent codebook from being updated during training
- **kmeans_init** (bool) - Optional - Default: False - Initialize codebook with k-means clustering
- **kmeans_iters** (int) - Optional - Default: 10 - Number of k-means iterations during initialization
- **sync_kmeans** (bool) - Optional - Default: True - Synchronize k-means across distributed processes
- **use_cosine_sim** (bool) - Optional - Default: False - Use cosine similarity distance instead of L2
- **layernorm_after_project_in** (bool) - Optional - Default: False - Apply layer normalization after input projection
- **threshold_ema_dead_code** (int) - Optional - Default: 0 - Replace codes used fewer than this many times
- **channel_last** (bool) - Optional - Default: True - Input format: True = (batch, seq, dim), False = (batch, dim, seq)
- **accept_image_fmap** (bool) - Optional - Default: False - Accept image format input (batch, channels, height, width)
- **accept_3d_fmap** (bool) - Optional - Default: False - Accept 3D feature map input (batch, channels, depth, height, width)
- **commitment_weight** (float) - Optional - Default: 1.0 - Weight of commitment loss in total loss
- **commitment_use_cross_entropy_loss** (bool) - Optional - Default: False - Use cross-entropy instead of MSE for commitment loss
- **orthogonal_reg_weight** (float) - Optional - Default: 0.0 - Weight of orthogonal regularization loss
- **orthogonal_reg_active_codes_only** (bool) - Optional - Default: False - Apply orthogonal loss only to used codes
- **orthogonal_reg_max_codes** (int | None) - Optional - Default: None - Limit orthogonal loss to top-N used codes
- **codebook_diversity_loss_weight** (float) - Optional - Default: 0.0 - Weight for codebook diversity loss
- **codebook_diversity_temperature** (float) - Optional - Default: 100.0 - Temperature for diversity loss softmax
- **stochastic_sample_codes** (bool) - Optional - Default: False - Stochastically sample codes instead of argmax
- **sample_codebook_temp** (float) - Optional - Default: 1.0 - Temperature for stochastic sampling
- **straight_through** (bool) - Optional - Default: False - Use straight-through estimator for gradients
- **rotation_trick** (bool | None) - Optional - Default: None - Use rotation trick for gradient propagation (2410.06424)
- **directional_reparam** (bool) - Optional - Default: False - Use directional reparameterization for learning
- **directional_reparam_variance** (float) - Optional - Default: 5e-3 - Noise variance for directional reparameterization
- **sync_codebook** (bool | None) - Optional - Default: None - Synchronize codebook updates across distributed processes
- **sync_affine_param** (bool) - Optional - Default: False - Synchronize affine parameters in distributed training
- **ema_update** (bool | None) - Optional - Default: None - Use EMA updates (auto-determined from other params)
- **vq_bridge** (nn.Module | None) - Optional - Default: None - Optional neural network to transform codebook
- **manual_ema_update** (bool) - Optional - Default: False - Manually control when EMA updates occur
- **learnable_codebook** (bool | None) - Optional - Default: None - Make codebook learnable via optimizer (auto-determined)
- **in_place_codebook_optimizer** (Callable) - Optional - Default: None - Custom optimizer for in-place codebook updates
- **manual_in_place_optimizer_update** (bool) - Optional - Default: False - Manually control optimizer updates
- **affine_param** (bool) - Optional - Default: False - Learn affine parameters for normalization
- **affine_param_batch_decay** (float) - Optional - Default: 0.99 - EMA decay for batch statistics
- **affine_param_codebook_decay** (float) - Optional - Default: 0.9 - EMA decay for codebook statistics
- **sync_update_v** (float) - Optional - Default: 0.0 - Synchronization weight (0 to 1) from VQTorch paper
- **return_zeros_for_masked_padding** (bool) - Optional - Default: True - Zero out masked positions in output
```

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

### Initialize and Use VectorQuantize

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Initialize the VectorQuantize module with specified dimensions and codebook size. Use it to quantize input tensors.

```python
import torch
from vector_quantize_pytorch import VectorQuantize

vq = VectorQuantize(
    dim = 256,
    codebook_size = 512,     # codebook size
    decay = 0.8,             # the exponential moving average decay, lower means the dictionary will change faster
    commitment_weight = 1.   # the weight on the commitment loss
)

x = torch.randn(1, 1024, 256)
quantized, indices, commit_loss = vq(x) # (1, 1024, 256), (1, 1024), (1)

```

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

### Finite Scalar Quantization (FSQ) Initialization

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Initialize a Finite Scalar Quantizer. This method simplifies vector quantization by rounding scalars, eliminating the need for commitment losses or EMA updates.

```python
import torch
from vector_quantize_pytorch import FSQ

quantizer = FSQ(
    levels = [8, 5, 5, 5]
)

x = torch.randn(1, 1024, 4) # 4 since there are 4 levels
xhat, indices = quantizer(x)

# (1, 1024, 4), (1, 1024)

assert torch.all(xhat == quantizer.indices_to_codes(indices))
```

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

### FSQ Class Initialization

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/fsq.md

Initializes the FSQ module with specified quantization levels and other parameters.

```APIDOC
## Class: FSQ

### Description
Initializes the Finite Scalar Quantization module.

### Constructor Parameters

#### levels
- **Type**: List[int] | Tuple[int, ...]
- **Description**: Quantization levels per dimension. The length of this list defines the codebook dimension.

#### dim
- **Type**: int | None
- **Default**: None
- **Description**: Input feature dimension. If None, it is inferred from `levels`.

#### num_codebooks
- **Type**: int
- **Default**: 1
- **Description**: Number of independent FSQ codebooks to use.

#### keep_num_codebooks_dim
- **Type**: bool | None
- **Default**: None
- **Description**: If True, the codebook dimension is kept in the output tensor.

#### scale
- **Type**: float | None
- **Default**: None
- **Description**: Scale factor to apply to the input and output.

#### allowed_dtypes
- **Type**: Tuple[torch.dtype, ...]
- **Default**: (torch.float32, torch.float64)
- **Description**: Tuple of allowed input data types for quantization.

#### channel_first
- **Type**: bool
- **Default**: False
- **Description**: If True, the input tensor is expected to be in channel-first format (batch, dim, seq). Otherwise, it's batch, seq, dim.

#### projection_has_bias
- **Type**: bool
- **Default**: True
- **Description**: Whether to include a bias term in the input and output projection layers.

#### return_indices
- **Type**: bool
- **Default**: True
- **Description**: If True, the output will include the quantization indices.

#### force_quantization_f32
- **Type**: bool
- **Default**: True
- **Description**: Forces the quantization computation to be performed in float32, regardless of the input dtype.

#### preserve_symmetry
- **Type**: bool
- **Default**: False
- **Description**: Ensures that the quantization ranges are symmetric.

#### noise_dropout
- **Type**: float
- **Default**: 0.0
- **Description**: Probability of applying dropout to the quantization noise.

#### bound_hard_clamp
- **Type**: bool
- **Default**: False
- **Description**: If True, the input values are hard-clamped to the quantization bounds.

#### orthogonal_rotation
- **Type**: bool
- **Default**: False
- **Description**: If True, applies an orthogonal rotation to the embeddings for potentially better utilization.
```

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

### Configure Simple Symmetric FSQ

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/fsq.md

Sets up a basic FSQ instance with 8 dimensions, each having 8 levels, and enables symmetry preservation. This configuration is useful for standard quantization tasks.

```python
fsq = FSQ(
    levels=[8] * 8,  # 8 dimensions, 8 levels each
    dim=512,
    preserve_symmetry=True
)
```

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

### Simple FSQ (No Learning)

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/README.md

Implement FSQ for deterministic quantization without learning, suitable when EMA updates are not desired. Define quantization levels, dimension, and symmetry preservation.

```python
from vector_quantize_pytorch import FSQ

fsq = FSQ(
    levels=[8, 6, 5],  # 240 total codes
    dim=768,
    preserve_symmetry=True
)

x = torch.randn(4, 1024, 768)
quantized, indices = fsq(x)

# Deterministic quantization, no EMA updates
```

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

### Handle Sequences and Video Features with LatentQuantize

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Demonstrates using LatentQuantize with sequence data (batch, seq, feat) and video data (batch, feat, time, height, width). The module adapts to these different input dimensionalities.

```python
import torch
from vector_quantize_pytorch import LatentQuantize

quantizer = LatentQuantize(
    levels = [5, 5, 8],
    dim = 16,
    commitment_loss_weight=0.1,
    quantization_loss_weight=0.1,
)

seq = torch.randn(1, 32, 16)
quantized, *_ = quantizer(seq)

# (1, 32, 16)

video_feats = torch.randn(1, 16, 10, 32, 32)
quantized, *_ = quantizer(video_feats)

# (1, 16, 10, 32, 32)
```

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

### Initialize RandomProjectionQuantizer

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/utilities.md

Initializes a RandomProjectionQuantizer for codebook initialization using random projections. It takes the input dimension, codebook size, and additional VectorQuantize arguments.

```python
from vector_quantize_pytorch import RandomProjectionQuantizer

rpq = RandomProjectionQuantizer(
    dim=256,
    codebook_size=512
)

x = torch.randn(4, 1024, 256)
quantized, indices = rpq(x)
```

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

### K-means Initialization

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/vector-quantize.md

Initialize the VectorQuantize module using K-means clustering. This is important for distributed training to ensure consistent initialization.

```python
vq = VectorQuantize(
    dim=256,
    codebook_size=512,
    kmeans_init=True,
    kmeans_iters=10,
    sync_kmeans=True  # important for distributed training
)
```

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

### Residual FSQ Initialization

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Initialize a Residual Finite Scalar Quantizer for improved audio encoding. This implementation is an improvised version based on a suggested idea.

```python
import torch
from vector_quantize_pytorch import ResidualFSQ

residual_fsq = ResidualFSQ(
    dim = 256,
    levels = [8, 5, 5, 3],
    num_quantizers = 8
)

x = torch.randn(1, 1024, 256)

residual_fsq.eval()

quantized, indices = residual_fsq(x)

# (1, 1024, 256), (1, 1024, 8)

quantized_out = residual_fsq.get_output_from_indices(indices)

# (1, 1024, 256)

assert torch.all(quantized == quantized_out)
```

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

### Initialize ResidualFSQ

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/residual-fsq.md

Instantiate ResidualFSQ with specified dimensions, quantization levels, and number of quantizers. The `levels` parameter can be a single integer for uniform levels across all dimensions, or a list/tuple for per-dimension or per-quantizer level specifications.

```python
from vector_quantize_pytorch import ResidualFSQ

rfsq = ResidualFSQ(
    dim=256,
    levels=[8, 6, 5],
    num_quantizers=4
)

x = torch.randn(4, 1024, 256)
quantized, indices, aux_loss = rfsq(x)

print(quantized.shape)  # torch.Size([4, 1024, 256])
print(indices.shape)    # torch.Size([4, 1024, 4])
```

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

### Initialize and Use ResidualVQ

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/README.md

Initialize the ResidualVQ module with a specified number of quantizers. Use it to recursively quantize residuals.

```python
import torch
from vector_quantize_pytorch import ResidualVQ

residual_vq = ResidualVQ(
    dim = 256,
    num_quantizers = 8,      # specify number of quantizers
    codebook_size = 1024,    # codebook size
)

x = torch.randn(1, 1024, 256)

quantized, indices, commit_loss = residual_vq(x)
print(quantized.shape, indices.shape, commit_loss.shape)
# (1, 1024, 256), (1, 1024, 8), (1, 8)

# if you need all the codes across the quantization layers, just pass return_all_codes = True

quantized, indices, commit_loss, all_codes = residual_vq(x, return_all_codes = True)

# (8, 1, 1024, 256)

```

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

### Configure Multi-Codebook FSQ

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/fsq.md

Initializes an FSQ instance with multiple independent quantizers. This allows for more complex quantization strategies by using several codebooks simultaneously.

```python
fsq = FSQ(
    levels=[16, 16, 16],
    dim=256,
    num_codebooks=4  # 4 independent quantizers
)
```

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

### HierarchicalVQ Configuration for Symmetric Hierarchy

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/hierarchical-vq.md

Shows how to initialize HierarchicalVQ for a symmetric hierarchy where the codebook size is consistent across layers, controlled by the `codebook_size` parameter for the first level and `num_layers`.

```python
hvq = HierarchicalVQ(
    dim=256,
    codebook_size=512,
    num_layers=3,
    commitment_weight=0.25
)
```

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

### Initialize FSQ

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/fsq.md

Initializes the FSQ module with specified levels, dimension, and number of codebooks. The codebook size is determined by the product of the levels.

```python
from vector_quantize_pytorch import FSQ

fsq = FSQ(
    levels=[8, 6, 5],  # 3 dimensions with different quantization levels
    dim=768,
    num_codebooks=1
)
```

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

### LFQ with Multiple Codebooks

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/lfq.md

Initialize LFQ to use multiple codebooks, specified by `num_codebooks`, for enhanced representation capacity.

```python
lfq = LFQ(
    dim=256,
    num_codebooks=4
)
```

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

### Initialize BinaryMapper

Source: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/_autodocs/api-reference/utilities.md

Initializes a BinaryMapper to convert continuous values to binary codes and vice versa. It supports a straight-through estimator and temperature for sampling.

```python
from vector_quantize_pytorch import BinaryMapper

mapper = BinaryMapper(num_bits=8)

x = torch.randn(4, 1024, 256)
quantized, codes = mapper(x)

print(quantized.shape)  # torch.Size([4, 1024, 256])
print(codes.shape)      # torch.Size([4, 1024, 8])
```