Spaces:

hugging-apps
/

greenwam

Running on Zero

greenwam / cosmos_framework /model /tokenizer /models /modules /quantizers /lfq.py

Initial GreenWAM demo

a1833c8 verified 21 days ago

16.8 kB

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""Lookup-Free Quantization (LFQ).
Proposed in https://arxiv.org/abs/2310.05737
In the simplest setup, each dimension is quantized into {-1, 1}.
An entropy penalty is used to encourage utilization.
References:
    - https://github.com/lucidrains/vector-quantize-pytorch
    - https://github.com/theAdamColton/ijepa-enhanced
"""

from __future__ import annotations

from collections import namedtuple
from math import ceil, log2
from typing import TYPE_CHECKING

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, reduce
from torch.nn import Module

if TYPE_CHECKING:
    from cosmos_framework.model.tokenizer.models.modules.sparse_tensor import SparseTensor

__all__ = [
    "LFQ",
    "LossBreakdown",
]


LossBreakdown = namedtuple(
    "LossBreakdown",
    ["per_sample_entropy", "codebook_entropy", "commitment", "avg_probs"],
)

_MAX_DIM_ONLY_CODEBOOK_BITS = 20


# Helper functions


def exists(v):
    """Check if value exists (is not None)."""
    return v is not None


def default(*args):
    """Return first non-None value from args (calling if callable)."""
    for arg in args:
        if exists(arg):
            return arg() if callable(arg) else arg
    return None


def entropy(prob: torch.Tensor) -> torch.Tensor:
    """Compute entropy of probability distribution."""
    return (-prob * torch.log(prob + 1e-5)).sum(dim=-1)


def mult_along_first_dims(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Multiply x by y along leading dimensions of y."""
    ndim_to_expand = x.ndim - y.ndim
    for _ in range(ndim_to_expand):
        y = y.unsqueeze(-1)
    return x * y


def masked_mean(x: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
    """Take mean of x elements not masked by m.
    The mean is taken along the shared leading dims of m.
    This is faster for torch-compile on batches than tensor indexing.
    """
    x = mult_along_first_dims(x, m)
    x = x / m.sum()
    return x.sum(tuple(range(m.ndim)))


def entropy_loss(
    logits: torch.Tensor,
    mask: torch.Tensor | None = None,
    temperature: float = 0.01,
    sample_minimization_weight: float = 1.0,
    batch_maximization_weight: float = 1.0,
    eps: float = 1e-5,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Compute entropy loss for LFQ.
    The loss encourages:
        - LOW per-sample entropy (focused code usage per sample)
        - HIGH batch-level entropy (diverse code usage across batch)
    Formula: loss = sample_entropy - batch_entropy
    NEGATIVE loss is GOOD: means batch entropy > sample entropy,
    indicating good codebook utilization with focused per-sample usage.
    Reference:
        https://github.com/google-research/magvit/
        LANGUAGE MODEL BEATS DIFFUSION — TOKENIZER IS KEY TO VISUAL GENERATION (2024)
    Args:
        logits: Affinities over the last dimension.
        mask: Optional mask for selective processing.
        temperature: Softmax temperature.
        sample_minimization_weight: Weight for per-sample entropy term.
        batch_maximization_weight: Weight for batch entropy term.
        eps: Small epsilon for numerical stability.
    Returns:
        Tuple of (per_sample_entropy, batch_entropy, loss).
    """
    probs = F.softmax(logits / temperature, -1)
    log_probs = F.log_softmax(logits / temperature + eps, -1)

    if mask is not None:
        avg_probs = masked_mean(probs, mask)
    else:
        avg_probs = reduce(probs, "... D -> D", "mean")

    avg_entropy = -torch.sum(avg_probs * torch.log(avg_probs + eps))

    sample_entropy = -torch.sum(probs * log_probs, -1)
    if mask is not None:
        sample_entropy = masked_mean(sample_entropy, mask).mean()
    else:
        sample_entropy = torch.mean(sample_entropy)

    # Key insight: loss = sample_entropy - batch_entropy
    # Negative loss is DESIRABLE (batch entropy > sample entropy)
    loss = (sample_minimization_weight * sample_entropy) - (batch_maximization_weight * avg_entropy)

    return sample_entropy, avg_entropy, loss


class LFQ(Module):
    """Lookup-Free Quantization module.
    Quantizes continuous values to binary codes {-1, 1} per dimension,
    with entropy loss to encourage codebook utilization.
    """

    def __init__(
        self,
        *,
        dim: int | None = None,
        codebook_size: int | None = None,
        num_codebooks: int = 1,
        sample_minimization_weight: float = 1.0,
        batch_maximization_weight: float = 1.0,
        token_factorization: bool = False,
        factorized_bits: list[int] = [9, 9],
    ) -> None:
        """Initialize LFQ.
        Args:
            dim: Input feature dimension (defaults to log2(codebook_size)).
            codebook_size: Size of codebook (must be power of 2).
            num_codebooks: Number of codebooks to use.
            sample_minimization_weight: Weight for per-sample entropy in loss.
            batch_maximization_weight: Weight for batch entropy in loss.
            token_factorization: Whether to use factorized tokens.
            factorized_bits: Bit split for factorized tokens.
        """
        super().__init__()

        if not exists(dim) and not exists(codebook_size):
            raise ValueError("Either dim or codebook_size must be specified for LFQ.")

        if codebook_size is None:
            if dim is None:
                raise ValueError("dim must be specified when codebook_size is omitted.")
            if dim > _MAX_DIM_ONLY_CODEBOOK_BITS:
                raise ValueError(
                    "LFQ dim-only construction materializes a 2**dim codebook; "
                    f"got dim={dim}. Pass codebook_size explicitly or use dim <= {_MAX_DIM_ONLY_CODEBOOK_BITS}."
                )
            resolved_codebook_size = 2**dim
        else:
            resolved_codebook_size = int(codebook_size)
        if not log2(resolved_codebook_size).is_integer():
            raise ValueError(f"codebook size must be power of 2 (suggested {2 ** ceil(log2(resolved_codebook_size))})")

        self.codebook_size = resolved_codebook_size
        self.codebook_dim = int(log2(self.codebook_size))

        codebook_dims = self.codebook_dim * num_codebooks
        dim = default(dim, codebook_dims)

        has_projections = dim != codebook_dims
        self.has_projections = has_projections

        self.dim = dim
        self.codebook_dim = self.codebook_dim
        self.num_codebooks = num_codebooks
        self.project_in = nn.Linear(self.dim, codebook_dims) if has_projections else nn.Identity()
        self.project_out = nn.Linear(codebook_dims, self.dim) if has_projections else nn.Identity()

        # For entropy loss
        self.sample_minimization_weight = sample_minimization_weight
        self.batch_maximization_weight = batch_maximization_weight

        # For token factorization
        self.token_factorization = token_factorization
        if not self.token_factorization:
            self.register_buffer("mask", 2 ** torch.arange(self.codebook_dim), persistent=False)
        else:
            self.factorized_bits = factorized_bits
            self.register_buffer("pre_mask", 2 ** torch.arange(factorized_bits[0]), persistent=False)
            self.register_buffer("post_mask", 2 ** torch.arange(factorized_bits[1]), persistent=False)

        self.register_buffer("zero", torch.tensor(0.0), persistent=False)

        # Build codebook
        all_codes = torch.arange(self.codebook_size)  # [K]
        bits = self.indices_to_bits(all_codes)  # [K,Z]
        codebook = bits * 2.0 - 1.0  # [K,Z]

        self.register_buffer("codebook", codebook, persistent=False)

    @property
    def dtype(self) -> torch.dtype:
        """Return dtype of codebook."""
        return self.codebook.dtype

    def indices_to_bits(self, x: torch.Tensor) -> torch.Tensor:
        """Convert indices to big-endian bits.
        Args:
            x: Long tensor of indices.
        Returns:
            Boolean tensor of bits.
        """
        mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long)
        x = (x.unsqueeze(-1) & mask) != 0
        return x

    def get_codebook_entry(
        self,
        x: torch.Tensor,
        bhwc: tuple[int, int, int, int],
        order: str,
    ) -> torch.Tensor:
        """Get codebook entry for given indices.
        Args:
            x: Index tensor.
            bhwc: Tuple of (batch, height, width, channels).
            order: 'pre' or 'post' for factorized tokens.
        Returns:
            Decoded tensor.
        """
        if self.token_factorization:
            if order == "pre":
                mask = 2 ** torch.arange(self.factorized_bits[0], device=x.device, dtype=torch.long)
            else:
                mask = 2 ** torch.arange(self.factorized_bits[1], device=x.device, dtype=torch.long)
        else:
            mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long)

        x = (x.unsqueeze(-1) & mask) != 0
        x = x * 2.0 - 1.0  # back to float
        b, h, w, c = bhwc
        x = rearrange(x, "b (h w) c -> b h w c", h=h, w=w, c=c)
        x = rearrange(x, "b h w c -> b c h w")
        return x

    def bits_to_indices(self, bits: torch.Tensor) -> torch.Tensor:
        """Convert big-endian bits to indices.
        Args:
            bits: Boolean tensor with bit dimension last.
        Returns:
            Long integer indices from 0 to codebook_size.
        """
        assert bits.shape[-1] == self.codebook_dim
        indices = 2 ** torch.arange(
            0,
            self.codebook_dim,
            1,
            dtype=torch.long,
            device=bits.device,
        )
        return (bits * indices).sum(-1)

    def decode(self, x: torch.Tensor) -> torch.Tensor:
        """Decode indices to continuous values.
        Args:
            x: Long tensor of codebook indices (..., NH) where NH is num_codebooks.
        Returns:
            Decoded tensor with values in {-1, 1}.
        """
        x = self.indices_to_bits(x)  # [...,NC,Z]
        x = x.to(self.dtype)  # [...,NC,Z]
        x = x * 2 - 1  # [...,NC,Z]
        x = rearrange(x, "... NC Z-> ... (NC Z)")  # [...,Dq]
        return self.project_out.to(x.dtype)(x)  # [...,D]

    def forward(
        self,
        x: "SparseTensor",
        inv_temperature: float = 100.0,
        return_loss_breakdown: bool = False,
        mask: torch.Tensor | None = None,
        return_loss: bool = True,
        fp32_loss_computation: bool = False,
    ) -> (
        tuple[torch.Tensor, torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor]
        | tuple[tuple[torch.Tensor, torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor], LossBreakdown]
    ):
        """Forward pass for LFQ on SparseTensor.
        Args:
            x: SparseTensor input with .feats, .coords, .shape, .layout.
            inv_temperature: Inverse temperature for entropy loss softmax.
            return_loss_breakdown: Whether to return detailed loss components.
            mask: Optional mask for selective processing.
            return_loss: Whether to compute training losses.
            fp32_loss_computation: Whether to compute losses in fp32.
        Returns:
            Tuple of (quantized_feats, indices, quantizer_loss).
            If return_loss_breakdown, also returns LossBreakdown.
        """
        # Extract features from sparse tensor
        N, feature_dim = x.shape

        expected_dim = self.num_codebooks * self.codebook_dim
        if feature_dim != self.dim:
            raise ValueError(f"Feature dimension {feature_dim} doesn't match LFQ input dimension {self.dim}.")

        features = self.project_in.to(x.dtype)(x.view(N, feature_dim))  # [N,Dq]
        features_reshaped = features.view(N, self.num_codebooks, self.codebook_dim)  # [N,NC,Z]

        # Quantization step
        codebook_value = torch.tensor(1.0, device=x.device, dtype=x.dtype)  # []
        quantized_values = torch.where(features_reshaped > 0, codebook_value, -codebook_value)  # [N,NC,Z]

        # Index calculation
        if self.token_factorization:
            pre_bits = quantized_values[..., : self.factorized_bits[0]]  # [N,NC,Zpre]
            post_bits = quantized_values[..., self.factorized_bits[0] :]  # [N,NC,Zpost]

            indices_pre = ((pre_bits > 0).int() * self.pre_mask.int()).sum(-1)  # [N,NC]
            indices_post = ((post_bits > 0).int() * self.post_mask.int()).sum(-1)  # [N,NC]

            indices_pre_flat = indices_pre.flatten()  # [N*NC]
            indices_post_flat = indices_post.flatten()  # [N*NC]
            sparse_indices_quantized = (indices_pre_flat, indices_post_flat)
        else:
            indices = ((quantized_values > 0).int() * self.mask.int()).sum(-1)  # [N,NC]
            sparse_indices_quantized = indices.flatten()  # [N*NC]

        # Entropy loss (training only)
        if self.training and return_loss:
            if fp32_loss_computation:
                features_flat_fp32 = features_reshaped.view(-1, self.codebook_dim).float()  # [N*NC,Z]
                codebook_fp32 = self.codebook.float()  # [K,Z]
            else:
                features_flat_fp32 = features_reshaped.view(-1, self.codebook_dim)  # [N*NC,Z]
                codebook_fp32 = self.codebook.to(features_flat_fp32.dtype)  # [K,Z]

            logits = 2 * torch.mm(features_flat_fp32, codebook_fp32.T)  # [N*NC,K]

            if mask is not None:
                if mask.shape[0] != N:
                    raise ValueError(f"Mask shape {mask.shape} doesn't match number of features {N}")
                mask_expanded = mask.unsqueeze(1).repeat(1, self.num_codebooks).view(-1)  # [N*NC]
            else:
                mask_expanded = None

            temperature = 1.0 / inv_temperature if inv_temperature > 0 else 0.01
            per_sample_entropy, codebook_entropy, entropy_aux_loss = entropy_loss(
                logits=logits,
                mask=mask_expanded,
                temperature=temperature,
                sample_minimization_weight=self.sample_minimization_weight,
                batch_maximization_weight=self.batch_maximization_weight,
            )
        else:
            dtype = torch.float32 if fp32_loss_computation else x.dtype
            per_sample_entropy = torch.tensor(0.0, dtype=dtype, device=x.device)  # []
            codebook_entropy = torch.tensor(0.0, dtype=dtype, device=x.device)  # []
            entropy_aux_loss = torch.tensor(0.0, dtype=dtype, device=x.device)  # []

        # Commitment loss
        if self.training:
            if fp32_loss_computation:
                features_fp32 = features_reshaped.float()  # [N,NC,Z]
                quantized_fp32 = quantized_values.float()  # [N,NC,Z]
            else:
                features_fp32 = features_reshaped  # [N,NC,Z]
                quantized_fp32 = quantized_values  # [N,NC,Z]

            commit_loss = F.mse_loss(features_fp32, quantized_fp32.detach(), reduction="none")  # [N,NC,Z]

            if mask is not None:
                mask_expanded = mask.view(N, 1, 1).expand_as(commit_loss)  # [N,NC,Z]
                commit_loss = commit_loss[mask_expanded].mean()  # []
            else:
                commit_loss = commit_loss.mean()  # []
        else:
            dtype = torch.float32 if fp32_loss_computation else x.dtype
            commit_loss = torch.tensor(0.0, dtype=dtype, device=x.device)  # []

        # Straight-through estimator
        quantized_values_ste = features_reshaped + (quantized_values - features_reshaped).detach()  # [N,NC,Z]

        # Output construction
        quantized_feats = quantized_values_ste.view(N, expected_dim)  # [N,Dq]
        quantized_feats = self.project_out.to(x.dtype)(quantized_feats)  # [N,D]

        # Ensure fp32 losses if requested
        if self.training and return_loss and fp32_loss_computation:
            entropy_aux_loss = entropy_aux_loss.float()  # []
            per_sample_entropy = per_sample_entropy.float()  # []
            codebook_entropy = codebook_entropy.float()  # []
            commit_loss = commit_loss.float()  # []

        quantizer_loss = commit_loss + entropy_aux_loss  # []
        ret = (quantized_feats, sparse_indices_quantized, quantizer_loss)

        if not return_loss_breakdown:
            return ret

        placeholder_dtype = torch.float32 if fp32_loss_computation else x.dtype
        return ret, LossBreakdown(
            per_sample_entropy,
            codebook_entropy,
            commit_loss,
            torch.tensor(0.0, dtype=placeholder_dtype, device=x.device),  # []
        )

	# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	# SPDX-License-Identifier: OpenMDW-1.1

	"""Lookup-Free Quantization (LFQ).

	Proposed in https://arxiv.org/abs/2310.05737

	In the simplest setup, each dimension is quantized into {-1, 1}.
	An entropy penalty is used to encourage utilization.

	References:
	- https://github.com/lucidrains/vector-quantize-pytorch
	- https://github.com/theAdamColton/ijepa-enhanced
	"""

	from __future__ import annotations

	from collections import namedtuple
	from math import ceil, log2
	from typing import TYPE_CHECKING

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from einops import rearrange, reduce
	from torch.nn import Module

	if TYPE_CHECKING:
	from cosmos_framework.model.tokenizer.models.modules.sparse_tensor import SparseTensor

	__all__ = [
	"LFQ",
	"LossBreakdown",
	]


	LossBreakdown = namedtuple(
	"LossBreakdown",
	["per_sample_entropy", "codebook_entropy", "commitment", "avg_probs"],
	)

	_MAX_DIM_ONLY_CODEBOOK_BITS = 20


	# Helper functions


	def exists(v):
	"""Check if value exists (is not None)."""
	return v is not None


	def default(*args):
	"""Return first non-None value from args (calling if callable)."""
	for arg in args:
	if exists(arg):
	return arg() if callable(arg) else arg
	return None


	def entropy(prob: torch.Tensor) -> torch.Tensor:
	"""Compute entropy of probability distribution."""
	return (-prob * torch.log(prob + 1e-5)).sum(dim=-1)


	def mult_along_first_dims(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
	"""Multiply x by y along leading dimensions of y."""
	ndim_to_expand = x.ndim - y.ndim
	for _ in range(ndim_to_expand):
	y = y.unsqueeze(-1)
	return x * y


	def masked_mean(x: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
	"""Take mean of x elements not masked by m.

	The mean is taken along the shared leading dims of m.
	This is faster for torch-compile on batches than tensor indexing.
	"""
	x = mult_along_first_dims(x, m)
	x = x / m.sum()
	return x.sum(tuple(range(m.ndim)))


	def entropy_loss(
	logits: torch.Tensor,
	mask: torch.Tensor \| None = None,
	temperature: float = 0.01,
	sample_minimization_weight: float = 1.0,
	batch_maximization_weight: float = 1.0,
	eps: float = 1e-5,
	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	"""Compute entropy loss for LFQ.

	The loss encourages:
	- LOW per-sample entropy (focused code usage per sample)
	- HIGH batch-level entropy (diverse code usage across batch)

	Formula: loss = sample_entropy - batch_entropy

	NEGATIVE loss is GOOD: means batch entropy > sample entropy,
	indicating good codebook utilization with focused per-sample usage.

	Reference:
	https://github.com/google-research/magvit/
	LANGUAGE MODEL BEATS DIFFUSION — TOKENIZER IS KEY TO VISUAL GENERATION (2024)

	Args:
	logits: Affinities over the last dimension.
	mask: Optional mask for selective processing.
	temperature: Softmax temperature.
	sample_minimization_weight: Weight for per-sample entropy term.
	batch_maximization_weight: Weight for batch entropy term.
	eps: Small epsilon for numerical stability.

	Returns:
	Tuple of (per_sample_entropy, batch_entropy, loss).
	"""
	probs = F.softmax(logits / temperature, -1)
	log_probs = F.log_softmax(logits / temperature + eps, -1)

	if mask is not None:
	avg_probs = masked_mean(probs, mask)
	else:
	avg_probs = reduce(probs, "... D -> D", "mean")

	avg_entropy = -torch.sum(avg_probs * torch.log(avg_probs + eps))

	sample_entropy = -torch.sum(probs * log_probs, -1)
	if mask is not None:
	sample_entropy = masked_mean(sample_entropy, mask).mean()
	else:
	sample_entropy = torch.mean(sample_entropy)

	# Key insight: loss = sample_entropy - batch_entropy
	# Negative loss is DESIRABLE (batch entropy > sample entropy)
	loss = (sample_minimization_weight * sample_entropy) - (batch_maximization_weight * avg_entropy)

	return sample_entropy, avg_entropy, loss


	class LFQ(Module):
	"""Lookup-Free Quantization module.

	Quantizes continuous values to binary codes {-1, 1} per dimension,
	with entropy loss to encourage codebook utilization.
	"""

	def __init__(
	self,
	*,
	dim: int \| None = None,
	codebook_size: int \| None = None,
	num_codebooks: int = 1,
	sample_minimization_weight: float = 1.0,
	batch_maximization_weight: float = 1.0,
	token_factorization: bool = False,
	factorized_bits: list[int] = [9, 9],
	) -> None:
	"""Initialize LFQ.

	Args:
	dim: Input feature dimension (defaults to log2(codebook_size)).
	codebook_size: Size of codebook (must be power of 2).
	num_codebooks: Number of codebooks to use.
	sample_minimization_weight: Weight for per-sample entropy in loss.
	batch_maximization_weight: Weight for batch entropy in loss.
	token_factorization: Whether to use factorized tokens.
	factorized_bits: Bit split for factorized tokens.
	"""
	super().__init__()

	if not exists(dim) and not exists(codebook_size):
	raise ValueError("Either dim or codebook_size must be specified for LFQ.")

	if codebook_size is None:
	if dim is None:
	raise ValueError("dim must be specified when codebook_size is omitted.")
	if dim > _MAX_DIM_ONLY_CODEBOOK_BITS:
	raise ValueError(
	"LFQ dim-only construction materializes a 2**dim codebook; "
	f"got dim={dim}. Pass codebook_size explicitly or use dim <= {_MAX_DIM_ONLY_CODEBOOK_BITS}."
	)
	resolved_codebook_size = 2**dim
	else:
	resolved_codebook_size = int(codebook_size)
	if not log2(resolved_codebook_size).is_integer():
	raise ValueError(f"codebook size must be power of 2 (suggested {2 ** ceil(log2(resolved_codebook_size))})")

	self.codebook_size = resolved_codebook_size
	self.codebook_dim = int(log2(self.codebook_size))

	codebook_dims = self.codebook_dim * num_codebooks
	dim = default(dim, codebook_dims)

	has_projections = dim != codebook_dims
	self.has_projections = has_projections

	self.dim = dim
	self.codebook_dim = self.codebook_dim
	self.num_codebooks = num_codebooks
	self.project_in = nn.Linear(self.dim, codebook_dims) if has_projections else nn.Identity()
	self.project_out = nn.Linear(codebook_dims, self.dim) if has_projections else nn.Identity()

	# For entropy loss
	self.sample_minimization_weight = sample_minimization_weight
	self.batch_maximization_weight = batch_maximization_weight

	# For token factorization
	self.token_factorization = token_factorization
	if not self.token_factorization:
	self.register_buffer("mask", 2 ** torch.arange(self.codebook_dim), persistent=False)
	else:
	self.factorized_bits = factorized_bits
	self.register_buffer("pre_mask", 2 ** torch.arange(factorized_bits[0]), persistent=False)
	self.register_buffer("post_mask", 2 ** torch.arange(factorized_bits[1]), persistent=False)

	self.register_buffer("zero", torch.tensor(0.0), persistent=False)

	# Build codebook
	all_codes = torch.arange(self.codebook_size) # [K]
	bits = self.indices_to_bits(all_codes) # [K,Z]
	codebook = bits * 2.0 - 1.0 # [K,Z]

	self.register_buffer("codebook", codebook, persistent=False)

	@property
	def dtype(self) -> torch.dtype:
	"""Return dtype of codebook."""
	return self.codebook.dtype

	def indices_to_bits(self, x: torch.Tensor) -> torch.Tensor:
	"""Convert indices to big-endian bits.

	Args:
	x: Long tensor of indices.

	Returns:
	Boolean tensor of bits.
	"""
	mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long)
	x = (x.unsqueeze(-1) & mask) != 0
	return x

	def get_codebook_entry(
	self,
	x: torch.Tensor,
	bhwc: tuple[int, int, int, int],
	order: str,
	) -> torch.Tensor:
	"""Get codebook entry for given indices.

	Args:
	x: Index tensor.
	bhwc: Tuple of (batch, height, width, channels).
	order: 'pre' or 'post' for factorized tokens.

	Returns:
	Decoded tensor.
	"""
	if self.token_factorization:
	if order == "pre":
	mask = 2 ** torch.arange(self.factorized_bits[0], device=x.device, dtype=torch.long)
	else:
	mask = 2 ** torch.arange(self.factorized_bits[1], device=x.device, dtype=torch.long)
	else:
	mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long)

	x = (x.unsqueeze(-1) & mask) != 0
	x = x * 2.0 - 1.0 # back to float
	b, h, w, c = bhwc
	x = rearrange(x, "b (h w) c -> b h w c", h=h, w=w, c=c)
	x = rearrange(x, "b h w c -> b c h w")
	return x

	def bits_to_indices(self, bits: torch.Tensor) -> torch.Tensor:
	"""Convert big-endian bits to indices.

	Args:
	bits: Boolean tensor with bit dimension last.

	Returns:
	Long integer indices from 0 to codebook_size.
	"""
	assert bits.shape[-1] == self.codebook_dim
	indices = 2 ** torch.arange(
	0,
	self.codebook_dim,
	1,
	dtype=torch.long,
	device=bits.device,
	)
	return (bits * indices).sum(-1)

	def decode(self, x: torch.Tensor) -> torch.Tensor:
	"""Decode indices to continuous values.

	Args:
	x: Long tensor of codebook indices (..., NH) where NH is num_codebooks.

	Returns:
	Decoded tensor with values in {-1, 1}.
	"""
	x = self.indices_to_bits(x) # [...,NC,Z]
	x = x.to(self.dtype) # [...,NC,Z]
	x = x * 2 - 1 # [...,NC,Z]
	x = rearrange(x, "... NC Z-> ... (NC Z)") # [...,Dq]
	return self.project_out.to(x.dtype)(x) # [...,D]

	def forward(
	self,
	x: "SparseTensor",
	inv_temperature: float = 100.0,
	return_loss_breakdown: bool = False,
	mask: torch.Tensor \| None = None,
	return_loss: bool = True,
	fp32_loss_computation: bool = False,
	) -> (
	tuple[torch.Tensor, torch.Tensor \| tuple[torch.Tensor, torch.Tensor], torch.Tensor]
	\| tuple[tuple[torch.Tensor, torch.Tensor \| tuple[torch.Tensor, torch.Tensor], torch.Tensor], LossBreakdown]
	):
	"""Forward pass for LFQ on SparseTensor.

	Args:
	x: SparseTensor input with .feats, .coords, .shape, .layout.
	inv_temperature: Inverse temperature for entropy loss softmax.
	return_loss_breakdown: Whether to return detailed loss components.
	mask: Optional mask for selective processing.
	return_loss: Whether to compute training losses.
	fp32_loss_computation: Whether to compute losses in fp32.

	Returns:
	Tuple of (quantized_feats, indices, quantizer_loss).
	If return_loss_breakdown, also returns LossBreakdown.
	"""
	# Extract features from sparse tensor
	N, feature_dim = x.shape

	expected_dim = self.num_codebooks * self.codebook_dim
	if feature_dim != self.dim:
	raise ValueError(f"Feature dimension {feature_dim} doesn't match LFQ input dimension {self.dim}.")

	features = self.project_in.to(x.dtype)(x.view(N, feature_dim)) # [N,Dq]
	features_reshaped = features.view(N, self.num_codebooks, self.codebook_dim) # [N,NC,Z]

	# Quantization step
	codebook_value = torch.tensor(1.0, device=x.device, dtype=x.dtype) # []
	quantized_values = torch.where(features_reshaped > 0, codebook_value, -codebook_value) # [N,NC,Z]

	# Index calculation
	if self.token_factorization:
	pre_bits = quantized_values[..., : self.factorized_bits[0]] # [N,NC,Zpre]
	post_bits = quantized_values[..., self.factorized_bits[0] :] # [N,NC,Zpost]

	indices_pre = ((pre_bits > 0).int() * self.pre_mask.int()).sum(-1) # [N,NC]
	indices_post = ((post_bits > 0).int() * self.post_mask.int()).sum(-1) # [N,NC]

	indices_pre_flat = indices_pre.flatten() # [N*NC]
	indices_post_flat = indices_post.flatten() # [N*NC]
	sparse_indices_quantized = (indices_pre_flat, indices_post_flat)
	else:
	indices = ((quantized_values > 0).int() * self.mask.int()).sum(-1) # [N,NC]
	sparse_indices_quantized = indices.flatten() # [N*NC]

	# Entropy loss (training only)
	if self.training and return_loss:
	if fp32_loss_computation:
	features_flat_fp32 = features_reshaped.view(-1, self.codebook_dim).float() # [N*NC,Z]
	codebook_fp32 = self.codebook.float() # [K,Z]
	else:
	features_flat_fp32 = features_reshaped.view(-1, self.codebook_dim) # [N*NC,Z]
	codebook_fp32 = self.codebook.to(features_flat_fp32.dtype) # [K,Z]

	logits = 2 * torch.mm(features_flat_fp32, codebook_fp32.T) # [N*NC,K]

	if mask is not None:
	if mask.shape[0] != N:
	raise ValueError(f"Mask shape {mask.shape} doesn't match number of features {N}")
	mask_expanded = mask.unsqueeze(1).repeat(1, self.num_codebooks).view(-1) # [N*NC]
	else:
	mask_expanded = None

	temperature = 1.0 / inv_temperature if inv_temperature > 0 else 0.01
	per_sample_entropy, codebook_entropy, entropy_aux_loss = entropy_loss(
	logits=logits,
	mask=mask_expanded,
	temperature=temperature,
	sample_minimization_weight=self.sample_minimization_weight,
	batch_maximization_weight=self.batch_maximization_weight,
	)
	else:
	dtype = torch.float32 if fp32_loss_computation else x.dtype
	per_sample_entropy = torch.tensor(0.0, dtype=dtype, device=x.device) # []
	codebook_entropy = torch.tensor(0.0, dtype=dtype, device=x.device) # []
	entropy_aux_loss = torch.tensor(0.0, dtype=dtype, device=x.device) # []

	# Commitment loss
	if self.training:
	if fp32_loss_computation:
	features_fp32 = features_reshaped.float() # [N,NC,Z]
	quantized_fp32 = quantized_values.float() # [N,NC,Z]
	else:
	features_fp32 = features_reshaped # [N,NC,Z]
	quantized_fp32 = quantized_values # [N,NC,Z]

	commit_loss = F.mse_loss(features_fp32, quantized_fp32.detach(), reduction="none") # [N,NC,Z]

	if mask is not None:
	mask_expanded = mask.view(N, 1, 1).expand_as(commit_loss) # [N,NC,Z]
	commit_loss = commit_loss[mask_expanded].mean() # []
	else:
	commit_loss = commit_loss.mean() # []
	else:
	dtype = torch.float32 if fp32_loss_computation else x.dtype
	commit_loss = torch.tensor(0.0, dtype=dtype, device=x.device) # []

	# Straight-through estimator
	quantized_values_ste = features_reshaped + (quantized_values - features_reshaped).detach() # [N,NC,Z]

	# Output construction
	quantized_feats = quantized_values_ste.view(N, expected_dim) # [N,Dq]
	quantized_feats = self.project_out.to(x.dtype)(quantized_feats) # [N,D]

	# Ensure fp32 losses if requested
	if self.training and return_loss and fp32_loss_computation:
	entropy_aux_loss = entropy_aux_loss.float() # []
	per_sample_entropy = per_sample_entropy.float() # []
	codebook_entropy = codebook_entropy.float() # []
	commit_loss = commit_loss.float() # []

	quantizer_loss = commit_loss + entropy_aux_loss # []
	ret = (quantized_feats, sparse_indices_quantized, quantizer_loss)

	if not return_loss_breakdown:
	return ret

	placeholder_dtype = torch.float32 if fp32_loss_computation else x.dtype
	return ret, LossBreakdown(
	per_sample_entropy,
	codebook_entropy,
	commit_loss,
	torch.tensor(0.0, dtype=placeholder_dtype, device=x.device), # []
	)