diveq 0.1.0__tar.gz

diveq-0.1.0/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026-present AaltoML
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

diveq-0.1.0/PKG-INFO

@@ -0,0 +1,146 @@
+Metadata-Version: 2.3
+Name: diveq
+Version: 0.1.0
+Summary: A PyTorch class for end-to-end training of vector quantization codebook in a deep neural network.
+Keywords: deep-learning,pytorch,vector-quantization,machine-learning,vq-vae,tokenization,representation-learning
+Author: Mohammad Hassan Vali
+Author-email: Mohammad Hassan Vali <mohammad.vali@aalto.fi>
+License: MIT License
+         
+         Copyright (c) 2026-present AaltoML
+         
+         Permission is hereby granted, free of charge, to any person obtaining a copy
+         of this software and associated documentation files (the "Software"), to deal
+         in the Software without restriction, including without limitation the rights
+         to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+         copies of the Software, and to permit persons to whom the Software is
+         furnished to do so, subject to the following conditions:
+         
+         The above copyright notice and this permission notice shall be included in all
+         copies or substantial portions of the Software.
+         
+         THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+         IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+         FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+         AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+         LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+         OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+         SOFTWARE.
+Requires-Dist: torch>=2.0.0
+Requires-Python: >=3.11
+Project-URL: Repository, https://github.com/AaltoML/DiVeQ_package
+Project-URL: Issues, https://github.com/AaltoML/DiVeQ_package/issues
+Project-URL: Paper, https://arxiv.org/pdf/2509.26469
+Description-Content-Type: text/markdown
+
+# Welcome to diveq
+`diveq` (short for differentiable vector quantization) is a tool for implementing and training vector quantization (VQ) in deep neural networks (DNNs), such as a VQ-VAE. It allows end-to-end training of DNNs that contain the non-differentiable VQ module, without any auxiliary losses and hyperparameter tunings. `diveq` is implemented via PyTorch, and it requires `python >= 3.11` and `torch >= 2.0.0` .
+
+![alt text](https://raw.githubusercontent.com/MHVali/DiVeQ_PyPI/main/example/diveq_teaser.png)
+
+`diveq` method is published as a research paper entitled [*"DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick"*](https://arxiv.org/pdf/2509.26469) in the International Conference on Learning Representations (ICLR) in 2026. You can find the original GitHub repository of the paper at [https://github.com/AaltoML/DiVeQ](https://github.com/AaltoML/DiVeQ).
+
+`diveq` package includes eight different vector quantization (VQ) methods:
+1. `from diveq import DIVEQ` optimizes the VQ codebook via DiVeQ technique. DiVeQ is the first proposed method in the paper that works as an ordinary VQ by mapping the input to codebook vectors.
+2. `from diveq import SFDIVEQ` optimizes the VQ codebook via Space-Filling DiVeQ technique. SF-DiVeQ is the second proposed method in the paper, different from ordinary VQ in a way that it maps the input to a space-filling curve constructed from codebook vectors.
+
+VQ variants that use multiple codebooks for vector quantization, i.e., Residual VQ and Product VQ:
+
+3. `from diveq import ResidualDIVEQ` optimizes the Residual VQ codebooks via DiVeQ technique.
+4. `from diveq import ResidualSFDIVEQ` optimizes the Residual VQ codebooks via SF-DiVeQ technique.
+5. `from diveq import ProductDIVEQ` optimizes the Product VQ codebooks via DiVeQ technique.
+6. `from diveq import ProductSFDIVEQ` optimizes the Product VQ codebooks via SF-DiVeQ technique.
+
+Variants of DiVeQ and SF-DiVeQ techniques that use deterministic quantization instead of stochastic quantization:
+
+7. `from diveq import DIVEQDetach` optimizes the VQ codebook via DiVeQ_Detach technique.
+8. `from diveq import SFDIVEQDetach` optimizes the VQ codebook via SF-DiVeQ_Detach technique.
+
+For more details on these eight different VQ methods, please see [the paper](https://arxiv.org/pdf/2509.26469).
+
+# Installation
+
+You can install `diveq` through `pip` by running:
+
+```bash
+pip install diveq
+```
+
+After installing `diveq`, you can verify the installation and package details by running:
+
+```bash
+python -m pip show diveq
+```
+
+# Usage Example
+
+Before using `diveq`, you have to install it using `pip install diveq`.
+
+Below you see a minimal example of how to import and use the `DIVEQ` optimization method as a vector quantizer in a model.
+
+```bash
+from diveq import DIVEQ
+vector_quantizer = DIVEQ(num_embeddings, embedding_dim)
+```
+
+- `vector_quantizer` is the vector quantization module that will be used for building the model.
+- `num_embeddings` and `embedding_dim` are the codebook size and dimension of each codebook entry, respectively. In the following, you can find the list of all parameters used in different vector quantization modules incorporated in `diveq` package.
+
+In the `example` directory of the [GitHub for `diveq` package](https://github.com/AaltoML/DiVeQ_package), we provide a code example of how vector quantization modules in `diveq` can be used in a vector quantized variational autoencoder (VQ-VAE). You can create the required environment to run the code by running:
+
+```bash
+cd example  #change directory to the example folder
+conda create --name diveq_example python=3.11
+conda activate diveq_example
+pip install -r requirements.txt
+```
+
+Then, you can train the VQ-VAE model by running:
+
+```bash
+python train.py
+```
+
+# List of Parameters
+Here, we provide the list of parameters that are used as inputs to eight different vector quantization methods included in `diveq` package.
+
+- `num_embeddings` (integer): Codebook size or the number of codewords in the codebook.
+- `embedding_dim` (integer): Dimensionality of each codebook entry or codeword.
+- `noise_var` (float): Variance of the directional noise for stochastic *DiVeQ*- and *SF-DiVeQ*-based methods.
+- `replacement_iters` (integer): Number of training iterations to apply codebook replacement.
+- `discard_threshold` (float): Threshold to discard the codebook entries that are used less than this threshold after *replacement_iters* iterations.
+- `perturb_eps` (float): Adjusts perturbation/shift magnitude from used codewords for codebook replacement.
+- `uniform_init` (bool): Whether to use uniform initialization. If False, codebook is initialized from a normal distribution.
+- `verbose` (bool): Whether to print codebook replacement status, i.e., to print how many unused codewords are replaced.
+- `skip_iters` (integer): Number of training iterations to skip quantization (for *SF-DiVeQ* and *SF-DiVeQ_Detach*) or to use *DiVeQ* quantization (for *Residual_SF-DiVeQ* and *Product_SF-DiVeQ*) in the custom initialization.
+- `avg_iters` (integer): Number of recent training iterations to extract latents for custom codebook initialization in Space-Filling Versions.
+- `latents_on_cpu` (bool): Whether to collect latents for custom initialization on CPU. If running out of CUDA memory, set it to True.
+- `allow_warning` (bool): Whether to print the warnings. The warnings will warn if the user inserts unusual values for the parameters.
+- `num_codebooks` (integer): Number of codebooks to be used for quantization in VQ variants of Residual VQ and Product VQ. All the codebooks will have the same size and dimensionality.
+
+# Important Notes about Parameters
+
+1. **Codebook Replacement:** Note that to prevent *codebook collapse*, we include a codebook replacement function (in cases where it is required) inside different quantization modules. Codebook replacement function is called after each `replacement_iters` training iterations, and it replaces the codewords which are used less than `discard_threshold` with perturbation of actively used codewords which are shifted by `perturb_eps` magnitude. If `verbose=True`, the status of how many unused codewords are replaced will be printed by the module. Note that the number of unused codewords should decrease over training, and it might take a while.
+
+2. **Variants of Vector Quantization:** Residual VQ and Product VQ are two variants of vector quantization, which are included in the `diveq` package. These variants utilize multiple codebooks for quantization, where `num_codebooks` determines the number of codebooks used in these VQ variants.
+
+3. **Space-Filling Methods:** Quantization methods based on Space-Filling (i.e., *SF-DiVeQ*, *SF-DiVeQ_Detach*, *Residual_SF-DiVeQ*, *Product_SF-DiVeQ*) use a custom initilization. *SF-DiVeQ* and *SF-DiVeQ_Detach* skip quantizing the latents for `skip_iters` training iterations, and initialize the codebook with an average of latents captured from `avg_iters` recent training iterations. After this custom initialization, they start to quantize the latents. *Residual_SF-DiVeQ* and *Product_SF-DiVeQ* work in the same way, but they apply *DiVeQ* for the first `skip_iters` training iterations. Note that if `avg_iters` value is set to a large value, CUDA might run out of memory, as there should be a large pull of latents to be stored for custom initialization. Therefore, the user can set `latents_on_cpu=True` to store the latents on CPU, or set a smaller value for `avg_iters`.
+
+4. **Detach Methods:** *DiVeQ_Detach* and *SF-DiVeQ_Detach* methods do not use directional noise. Therefore, they do not need to set the `noise_var` parameter.
+
+For further details about different vector quantization methods in the `diveq` package and their corresponding parameters, please see the details provided in the Python codes in `src` directory of the `diveq` package.
+
+# Citation
+If this package contributed to your work, please consider citing it:
+
+```
+@InProceedings{vali2026diveq,
+ title={{DiVeQ}: {D}ifferentiable {V}ector {Q}uantization {U}sing the {R}eparameterization {T}rick},
+ author={Vali, Mohammad Hassan and Bäckström, Tom and Solin, Arno},
+ booktitle={International Conference on Learning Representations (ICLR)},
+ year={2026}
+}
+```
+
+# License
+`diveq` was developed by <span property="cc:attributionName">Mohammad Hassan Vali</span>, part of the <a href="https://users.aalto.fi/~asolin/group/" target="_blank">AaltoML research group from Aalto University</a> and is licensed under MIT license.

diveq-0.1.0/README.md

@@ -0,0 +1,111 @@
+# Welcome to diveq
+`diveq` (short for differentiable vector quantization) is a tool for implementing and training vector quantization (VQ) in deep neural networks (DNNs), such as a VQ-VAE. It allows end-to-end training of DNNs that contain the non-differentiable VQ module, without any auxiliary losses and hyperparameter tunings. `diveq` is implemented via PyTorch, and it requires `python >= 3.11` and `torch >= 2.0.0` .
+
+![alt text](https://raw.githubusercontent.com/MHVali/DiVeQ_PyPI/main/example/diveq_teaser.png)
+
+`diveq` method is published as a research paper entitled [*"DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick"*](https://arxiv.org/pdf/2509.26469) in the International Conference on Learning Representations (ICLR) in 2026. You can find the original GitHub repository of the paper at [https://github.com/AaltoML/DiVeQ](https://github.com/AaltoML/DiVeQ).
+
+`diveq` package includes eight different vector quantization (VQ) methods:
+1. `from diveq import DIVEQ` optimizes the VQ codebook via DiVeQ technique. DiVeQ is the first proposed method in the paper that works as an ordinary VQ by mapping the input to codebook vectors.
+2. `from diveq import SFDIVEQ` optimizes the VQ codebook via Space-Filling DiVeQ technique. SF-DiVeQ is the second proposed method in the paper, different from ordinary VQ in a way that it maps the input to a space-filling curve constructed from codebook vectors.
+
+VQ variants that use multiple codebooks for vector quantization, i.e., Residual VQ and Product VQ:
+
+3. `from diveq import ResidualDIVEQ` optimizes the Residual VQ codebooks via DiVeQ technique.
+4. `from diveq import ResidualSFDIVEQ` optimizes the Residual VQ codebooks via SF-DiVeQ technique.
+5. `from diveq import ProductDIVEQ` optimizes the Product VQ codebooks via DiVeQ technique.
+6. `from diveq import ProductSFDIVEQ` optimizes the Product VQ codebooks via SF-DiVeQ technique.
+
+Variants of DiVeQ and SF-DiVeQ techniques that use deterministic quantization instead of stochastic quantization:
+
+7. `from diveq import DIVEQDetach` optimizes the VQ codebook via DiVeQ_Detach technique.
+8. `from diveq import SFDIVEQDetach` optimizes the VQ codebook via SF-DiVeQ_Detach technique.
+
+For more details on these eight different VQ methods, please see [the paper](https://arxiv.org/pdf/2509.26469).
+
+# Installation
+
+You can install `diveq` through `pip` by running:
+
+```bash
+pip install diveq
+```
+
+After installing `diveq`, you can verify the installation and package details by running:
+
+```bash
+python -m pip show diveq
+```
+
+# Usage Example
+
+Before using `diveq`, you have to install it using `pip install diveq`.
+
+Below you see a minimal example of how to import and use the `DIVEQ` optimization method as a vector quantizer in a model.
+
+```bash
+from diveq import DIVEQ
+vector_quantizer = DIVEQ(num_embeddings, embedding_dim)
+```
+
+- `vector_quantizer` is the vector quantization module that will be used for building the model.
+- `num_embeddings` and `embedding_dim` are the codebook size and dimension of each codebook entry, respectively. In the following, you can find the list of all parameters used in different vector quantization modules incorporated in `diveq` package.
+
+In the `example` directory of the [GitHub for `diveq` package](https://github.com/AaltoML/DiVeQ_package), we provide a code example of how vector quantization modules in `diveq` can be used in a vector quantized variational autoencoder (VQ-VAE). You can create the required environment to run the code by running:
+
+```bash
+cd example  #change directory to the example folder
+conda create --name diveq_example python=3.11
+conda activate diveq_example
+pip install -r requirements.txt
+```
+
+Then, you can train the VQ-VAE model by running:
+
+```bash
+python train.py
+```
+
+# List of Parameters
+Here, we provide the list of parameters that are used as inputs to eight different vector quantization methods included in `diveq` package.
+
+- `num_embeddings` (integer): Codebook size or the number of codewords in the codebook.
+- `embedding_dim` (integer): Dimensionality of each codebook entry or codeword.
+- `noise_var` (float): Variance of the directional noise for stochastic *DiVeQ*- and *SF-DiVeQ*-based methods.
+- `replacement_iters` (integer): Number of training iterations to apply codebook replacement.
+- `discard_threshold` (float): Threshold to discard the codebook entries that are used less than this threshold after *replacement_iters* iterations.
+- `perturb_eps` (float): Adjusts perturbation/shift magnitude from used codewords for codebook replacement.
+- `uniform_init` (bool): Whether to use uniform initialization. If False, codebook is initialized from a normal distribution.
+- `verbose` (bool): Whether to print codebook replacement status, i.e., to print how many unused codewords are replaced.
+- `skip_iters` (integer): Number of training iterations to skip quantization (for *SF-DiVeQ* and *SF-DiVeQ_Detach*) or to use *DiVeQ* quantization (for *Residual_SF-DiVeQ* and *Product_SF-DiVeQ*) in the custom initialization.
+- `avg_iters` (integer): Number of recent training iterations to extract latents for custom codebook initialization in Space-Filling Versions.
+- `latents_on_cpu` (bool): Whether to collect latents for custom initialization on CPU. If running out of CUDA memory, set it to True.
+- `allow_warning` (bool): Whether to print the warnings. The warnings will warn if the user inserts unusual values for the parameters.
+- `num_codebooks` (integer): Number of codebooks to be used for quantization in VQ variants of Residual VQ and Product VQ. All the codebooks will have the same size and dimensionality.
+
+# Important Notes about Parameters
+
+1. **Codebook Replacement:** Note that to prevent *codebook collapse*, we include a codebook replacement function (in cases where it is required) inside different quantization modules. Codebook replacement function is called after each `replacement_iters` training iterations, and it replaces the codewords which are used less than `discard_threshold` with perturbation of actively used codewords which are shifted by `perturb_eps` magnitude. If `verbose=True`, the status of how many unused codewords are replaced will be printed by the module. Note that the number of unused codewords should decrease over training, and it might take a while.
+
+2. **Variants of Vector Quantization:** Residual VQ and Product VQ are two variants of vector quantization, which are included in the `diveq` package. These variants utilize multiple codebooks for quantization, where `num_codebooks` determines the number of codebooks used in these VQ variants.
+
+3. **Space-Filling Methods:** Quantization methods based on Space-Filling (i.e., *SF-DiVeQ*, *SF-DiVeQ_Detach*, *Residual_SF-DiVeQ*, *Product_SF-DiVeQ*) use a custom initilization. *SF-DiVeQ* and *SF-DiVeQ_Detach* skip quantizing the latents for `skip_iters` training iterations, and initialize the codebook with an average of latents captured from `avg_iters` recent training iterations. After this custom initialization, they start to quantize the latents. *Residual_SF-DiVeQ* and *Product_SF-DiVeQ* work in the same way, but they apply *DiVeQ* for the first `skip_iters` training iterations. Note that if `avg_iters` value is set to a large value, CUDA might run out of memory, as there should be a large pull of latents to be stored for custom initialization. Therefore, the user can set `latents_on_cpu=True` to store the latents on CPU, or set a smaller value for `avg_iters`.
+
+4. **Detach Methods:** *DiVeQ_Detach* and *SF-DiVeQ_Detach* methods do not use directional noise. Therefore, they do not need to set the `noise_var` parameter.
+
+For further details about different vector quantization methods in the `diveq` package and their corresponding parameters, please see the details provided in the Python codes in `src` directory of the `diveq` package.
+
+# Citation
+If this package contributed to your work, please consider citing it:
+
+```
+@InProceedings{vali2026diveq,
+ title={{DiVeQ}: {D}ifferentiable {V}ector {Q}uantization {U}sing the {R}eparameterization {T}rick},
+ author={Vali, Mohammad Hassan and Bäckström, Tom and Solin, Arno},
+ booktitle={International Conference on Learning Representations (ICLR)},
+ year={2026}
+}
+```
+
+# License
+`diveq` was developed by <span property="cc:attributionName">Mohammad Hassan Vali</span>, part of the <a href="https://users.aalto.fi/~asolin/group/" target="_blank">AaltoML research group from Aalto University</a> and is licensed under MIT license.

diveq-0.1.0/pyproject.toml

@@ -0,0 +1,30 @@
+[project]
+name = "diveq"
+authors = [{name = "Mohammad Hassan Vali", email = "mohammad.vali@aalto.fi"}]
+version = "0.1.0"
+description = "A PyTorch class for end-to-end training of vector quantization codebook in a deep neural network."
+readme = "README.md"
+license = { file = "LICENSE.txt" }
+requires-python = ">=3.11"
+dependencies = [
+    "torch>=2.0.0",
+]
+
+keywords = [
+    "deep-learning",
+    "pytorch",
+    "vector-quantization",
+    "machine-learning",
+    "vq-vae",
+    "tokenization",
+    "representation-learning"
+]
+
+[project.urls]
+Repository = "https://github.com/AaltoML/DiVeQ_package"
+Issues = "https://github.com/AaltoML/DiVeQ_package/issues"
+Paper = "https://arxiv.org/pdf/2509.26469"
+
+[build-system]
+requires = ["uv_build>=0.9.10,<0.10.0"]
+build-backend = "uv_build"

diveq-0.1.0/src/diveq/__init__.py

@@ -0,0 +1,8 @@
+from .diveq import DIVEQ
+from .sf_diveq import SFDIVEQ
+from .diveq_detach import DIVEQDetach
+from .sf_diveq_detach import SFDIVEQDetach
+from .residual_diveq import ResidualDIVEQ
+from .residual_sf_diveq import ResidualSFDIVEQ
+from .product_diveq import ProductDIVEQ
+from .product_sf_diveq import ProductSFDIVEQ

diveq-0.1.0/src/diveq/diveq.py

@@ -0,0 +1,241 @@
+import torch
+import torch.nn as nn
+from torch.distributions import normal
+from typing import Tuple
+import warnings
+
+class DIVEQ(nn.Module):
+    """
+    DiVeQ: Differentiable Vector Quantization (VQ) module that allows end-to-end
+    training of VQ-based models without any auxiliary losses or hyperparameter tunings.
+    The module encompasses codebook replacement function which discards unused codebook
+    entries during training.
+
+    Args:
+        - num_embeddings (int): Codebook size (No. of codewords).
+        - embedding_dim (int): Dimensionality of embeddings.
+        - noise_var (float): Variance of the directional noise in DiVeQ.
+            Recommended noise_var < 1e-2.
+        - replacement_iters (int): Replacement interval (number of training iterations
+            to apply codebook replacement). Recommended 50 < replacement_iters < 300.
+        - discard_threshold (float): Threshold to discard the codebook entries that are
+            used less than this threshold after "replacement_iters" iterations.
+            Recommended 0.01 < discard_threshold < 0.05. discard_threshold must be in
+            the range of [0,1] such that discard_threshold=0.01 means to discard the
+            codebook entries which are used less than 1 percent.
+        - perturb_eps (float): Adjusts perturbation/shift magnitude from used codewords
+            for codebook replacement.
+        - uniform_init (bool): Whether to initialize codebook with uniform distribution.
+            If False, the codebook is initialized from a normal distribution.
+        - allow_warning (bool): Whether to print the warnings.
+        - verbose (bool): Whether to print codebook replacement status.
+
+    Returns:
+        - z_q (torch.Tensor): Differentiable quantized input/latent. shape (N, D)
+        - indices (torch.Tensor): Selected codebook indices. shape (N, )
+        - perplexity (float): Codebook perplexity (average codebook usage)
+    """
+    def __init__(
+            self,
+            num_embeddings: int,
+            embedding_dim: int,
+            noise_var: float = 0.001,
+            replacement_iters: int = 100,
+            discard_threshold: float = 0.01,
+            perturb_eps: float = 1e-9,
+            uniform_init: bool = True,
+            allow_warning: bool = True,
+            verbose: bool = True,
+    ):
+        super().__init__()
+
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.noise_var = noise_var
+        self.replacement_iters = replacement_iters
+        self.discard_threshold = discard_threshold
+        self.perturb_eps = perturb_eps
+        self.uniform_init = uniform_init
+        self.allow_warning = allow_warning
+        self.verbose = verbose
+
+        self._check_constraints()
+
+        # ---------------- User warnings ----------------
+        if allow_warning:
+            if noise_var > 0.01:
+                warnings.warn(f"`noise_var` is set to {noise_var}, which is"
+                              f" quite large. Values > 0.01 may overshoot"
+                              f" nearest-neighbor mapping.", UserWarning)
+            if replacement_iters < 50:
+                warnings.warn(f"`replacement_iters` is set to"
+                              f" {replacement_iters}, which is quite small. Values < 50"
+                              f" may cause too early and frequent codebook"
+                              f" replacements.", UserWarning)
+            elif replacement_iters > 300:
+                warnings.warn(f"`replacement_iters` is set to"
+                              f" {replacement_iters}, which is quite large."
+                              f" Values > 300 may cause too late and sporadic codebook"
+                              f" replacements.", UserWarning)
+
+            if discard_threshold > 0.05:
+                warnings.warn(f"`discard_threshold` is set to"
+                              f" {discard_threshold}, which is quite large."
+                              f" Values > 0.05 may discard a portion of suitable but"
+                              f" rarely used codewords.", UserWarning)
+
+            if perturb_eps > 1e-6:
+                warnings.warn(f"`perturb_eps` is set to {perturb_eps}, which"
+                              f" is quite large. Values > 1e-6 may cause big"
+                              f" perturbation/shift from used codewords.", UserWarning)
+
+        # ---------------- Codebook initialization ----------------
+        if uniform_init:
+            codebook = (torch.rand((self.num_embeddings, self.embedding_dim))
+                        * (1 / self.num_embeddings))
+        else:
+            codebook = (torch.randn((self.num_embeddings, self.embedding_dim))
+                        * (1 / self.num_embeddings))
+
+        self.codebook = torch.nn.Parameter(codebook, requires_grad=True)
+
+        # ---------------- Tensors used for codebook replacement ----------------
+        self.register_buffer("codebook_usage", torch.zeros(self.num_embeddings,
+                                                           dtype=torch.int32))
+        self.register_buffer("iter_counter", torch.zeros(1, dtype=torch.int32))
+
+    # ---------------- Forward pass (Core API) ----------------
+    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, float]:
+        """
+        Args:
+            - z (torch.Tensor): input/latent. shape (N, D)
+
+        Returns:
+            - z_q (torch.Tensor): Differentiable quantized input/latent. shape (N, D)
+            - indices (torch.Tensor): Selected codebook indices. shape (N, )
+            - perplexity (float): Codebook perplexity (average codebook usage).
+        """
+        self._check_input(z)
+
+        # Calculate distances
+        distances = (torch.sum(z.pow(2), dim=1, keepdim=True)
+                     + torch.sum(self.codebook.pow(2), dim=1)
+                     - 2 * torch.matmul(z, self.codebook.t()))
+
+        indices = torch.argmin(distances, dim=1)
+
+        z_hard_quantized = self.codebook[indices] # Non-differentiable quantized input
+
+        direction = z_hard_quantized - z
+        random_vectors = (normal.Normal(0, self.noise_var).sample(z.shape)
+                          .to(z.device) + direction)
+        normalized = random_vectors / torch.linalg.norm(random_vectors, dim=1,
+                                                        keepdim=True).clamp_min(1e-12)
+        error_magnitude = torch.linalg.norm(z_hard_quantized - z, dim=1, keepdim=True)
+
+        vq_error = error_magnitude * normalized.detach()
+        z_q = z + vq_error # Differentiable quantized input
+
+        # Perplexity Computation
+        perplexity = self._compute_perplexity(indices)
+
+        # Track used indices for codebook replacement
+        with torch.no_grad():
+            self.codebook_usage[indices] += 1
+            self.iter_counter += 1
+            if self.iter_counter.item() % self.replacement_iters == 0:
+                self._replace_unused_entries() # Applies codebook replacement
+
+        return z_q, indices, perplexity
+
+    # ---------------- Quantization for Inference ----------------
+    @torch.no_grad()
+    def inference(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, float]:
+        """
+        Deterministic hard quantization by mapping the input to the nearest codeword.
+        Args:
+            - z (torch.Tensor): input/latent. shape (N, D)
+
+        Returns:
+            - z_q_hard (torch.Tensor): Hard quantized input/latent. shape (N, D)
+            - indices (torch.Tensor): Selected codebook indices. shape (N, )
+            - perplexity (float): Codebook perplexity (average codebook usage).
+        """
+        self._check_input(z)
+        distances = (torch.sum(z.pow(2), dim=1, keepdim=True)
+                     + torch.sum(self.codebook.pow(2), dim=1)
+                     - 2.0 * torch.matmul(z, self.codebook.t()))
+        indices = torch.argmin(distances, dim=1)
+        perplexity = self._compute_perplexity(indices)
+        z_q_hard = self.codebook[indices]
+        return z_q_hard, indices, perplexity
+
+    # ---------------- Utility functions ----------------
+    def _check_constraints(self,) -> None:
+        if self.noise_var <= 0.0:
+            raise ValueError("`noise_var` must be a positive float value. To have more"
+                             " precise nearest-neighbor assignments, it is recommended"
+                             " that noise_var < 1e-2.")
+        if (self.replacement_iters <= 0) or (type(self.replacement_iters) is not int):
+            raise ValueError("`replacement_iters` must be a positive integer value."
+                             " It is recommended that 50 < replacement_iters < 300.")
+        if (self.discard_threshold < 0.0) or (self.discard_threshold > 1.0):
+            raise ValueError("`discard_threshold` must be in the range of [0,1]. It is"
+                             " recommended that 0.01 < discard_threshold < 0.05,"
+                             " such that discard_threshold=0.01 means to discard the"
+                             " codebook entries which are used less than 1 percent.")
+
+    def _check_input(self, z: torch.Tensor) -> None:
+        if z.ndim != 2:
+            raise ValueError("DiVeQ input must have the shape of (N, D), where N is"
+                             " the No. of input samples,and D is the embedding"
+                             " dimensionality.")
+        if z.size(1) != self.embedding_dim:
+            raise ValueError(f"DiVeQ input.shape[1] must match the embedding"
+                             f" dimensionality that is {self.embedding_dim}.")
+
+    def _compute_perplexity(self, indices: torch.Tensor) -> float:
+        encodings = torch.zeros(indices.shape[0], self.num_embeddings,
+                                device=indices.device)
+        encodings.scatter_(1, indices.unsqueeze(1), 1)
+        avg_probs = torch.mean(encodings, dim=0)
+        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
+        return perplexity.item()
+
+    def _replace_unused_entries(self) -> None:
+        with torch.no_grad():
+            usage_ratio = self.codebook_usage / self.replacement_iters
+            unused_indices = torch.where(usage_ratio < self.discard_threshold)[0]
+            used_indices = torch.where(usage_ratio >= self.discard_threshold)[0]
+
+            if unused_indices.numel() == 0 or used_indices.numel() == 0:
+                self.codebook_usage.zero_()
+                return
+
+            unused_count = unused_indices.numel()
+            used_probs = self.codebook_usage[used_indices] / torch.sum(
+                self.codebook_usage[used_indices])
+            randomly_sampled_indices = used_probs.multinomial(num_samples=unused_count,
+                                                              replacement=True)
+            sampled_indices = used_indices[randomly_sampled_indices]
+            used_codebooks = self.codebook[sampled_indices].clone()
+
+            self.codebook[unused_indices] = (used_codebooks +
+                         self.perturb_eps * torch.randn_like(used_codebooks)).clone()
+            self.codebook_usage.zero_()
+
+            if self.verbose:
+                print("\n***** Replaced " + str(unused_count) + " codewords *****")
+
+    def extra_repr(self) -> str:
+        return (
+            f"num_embeddings={self.num_embeddings}, "
+            f"embedding_dim={self.embedding_dim}, "
+            f"noise_var={self.noise_var}, "
+            f"replacement_iters={self.replacement_iters}, "
+            f"discard_threshold={self.discard_threshold}, "
+            f"perturb_eps={self.perturb_eps}, "
+            f"uniform_init={self.uniform_init}, "
+            f"allow_warning={self.allow_warning}, "
+            f"verbose={self.verbose}"
+        )