Tokenizers

All tokenizers subclass BaseTokenizer and share a common interface (encode / decode / tokenize / detokenize / reconstruct / forward / from_pretrained / save_pretrained).

class medtokenizers.BaseTokenizer(dim=None, name='BaseTokenizer')[source][source]

Bases: abc.ABC, torch.nn.modules.module.Module

Abstract base class for all tokenizers with HuggingFace Hub integration.

Tokenizers transform input images/volumes into latent representations (continuous or discrete) and back. This class provides:

  1. Abstract encode/decode interface - Subclasses implement specifics

  2. HuggingFace Hub integration - save_pretrained, from_pretrained, push_to_hub

  3. Batch processing utilities - For large datasets

  4. Sliding window inference - For volumes larger than GPU memory

Architecture Overview

A tokenizer consists of: - Encoder: Compresses input to low-dimensional latent - Quantizer (discrete only): Maps continuous latent to discrete codes - Decoder: Reconstructs input from latent

` Input -> Encoder -> [Quantizer] -> Latent -> Decoder -> Reconstruction `

HuggingFace Hub Integration

All tokenizers can be saved and loaded from the HuggingFace Hub:

>>> model.save_pretrained("./my-tokenizer")
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
>>> model.push_to_hub("username/my-tokenizer")

Thread Safety

Models should be used in single-threaded contexts or with appropriate synchronization. The eval()/train() mode switching is NOT thread-safe.

type dim:

Optional[int], default: None

param dim:

Spatial dimensionality (2 for images, 3 for volumes)

type name:

str, default: 'BaseTokenizer'

param name:

Human-readable model name

config_name: str = 'config.json'
weights_name: str = 'pytorch_model.bin'
__init__(dim=None, name='BaseTokenizer')[source][source]
Parameters:

  • dim (Optional[int], default: None)

  • name (str, default: 'BaseTokenizer')

config: dict[str, Any]
abstractmethod encode(x)[source][source]

Encode input to latent representation.

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input image/volume tensor

Return type:

tuple[Tensor, ...]

Returns:

Tuple containing latent representation and additional outputs (varies by subclass - e.g., KL divergence for VAE)

abstractmethod decode(z)[source][source]

Decode latent representation to output.

Parameters:

z (Float[Tensor, 'batch channels *spatial']) – Latent tensor (continuous codes or quantized codes)

Return type:

Float[Tensor, 'batch channels *spatial']

Returns:

Reconstructed output with same spatial shape as original input

abstractmethod forward(x)[source][source]

Full forward pass: encode -> [quantize] -> decode.

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input tensor

Returns:

  • β€˜reconstructions’: Reconstructed output

  • Additional keys vary by subclass (posteriors, quant_loss, etc.)

Return type:

Dictionary containing at minimum

inference_mode()[source][source]

Context manager for optimized inference.

Configures the model for maximum inference performance: - Sets model to eval mode - Disables gradient computation - Uses torch.inference_mode for additional optimizations

The model is restored to its previous state upon exit.

Return type:

Generator[BaseTokenizer, None, None]

Example

>>> with model.inference_mode():
...     latents = model.tokenize(volume)
...     recon = model.detokenize(latents)
Yields:

self – The model instance for method chaining

num_parameters()[source][source]

Get total number of learnable parameters.

Return type:

int

Returns:

Sum of numel() for all parameters

tokenize(x)[source][source]

Encode input to latent space (convenience method).

Subclasses should override with appropriate return type.

Parameters:

x (Tensor) – Input tensor

Return type:

Tensor

Returns:

Latent representation (continuous or discrete indices)

detokenize(z)[source][source]

Decode from latent space (convenience method).

Parameters:

z (Tensor) – Latent tensor

Return type:

Tensor

Returns:

Reconstructed output

compile(mode='reduce-overhead', fullgraph=False, **kwargs)[source][source]

Compile model with torch.compile() for faster inference.

Uses PyTorch 2.0+ compilation to optimize the model graph. The compiled model maintains the same interface but runs faster, especially for repeated inference calls.

Parameters:

  • mode (str, default: 'reduce-overhead') – Compilation mode. Options: - β€œreduce-overhead”: Best for small batches (default) - β€œmax-autotune”: Best throughput, longer warmup - β€œdefault”: Balanced compilation

  • fullgraph (bool, default: False) – If True, require full graph compilation (stricter)

  • **kwargs – Additional arguments passed to torch.compile()

Return type:

BaseTokenizer

Returns:

Compiled model (self, for method chaining)

Example

>>> model = ContinuousTokenizer.from_pretrained("path/to/model")
>>> model = model.compile(mode="reduce-overhead")
>>> # First call triggers compilation (slower)
>>> latents = model.tokenize(batch)
>>> # Subsequent calls are faster
>>> latents = model.tokenize(batch2)

Note

  • Compilation happens on first forward pass

  • Different input shapes may trigger recompilation

  • Use dynamic=False (default) for fixed input sizes

save_pretrained(save_directory, push_to_hub=False, **kwargs)[source][source]

Save model weights and configuration to directory.

Creates a directory structure compatible with from_pretrained(): ` save_directory/ β”œβ”€β”€ config.json      # Model configuration └── pytorch_model.bin # Model weights `

Parameters:

  • save_directory (str | Path) – Path to save model

  • push_to_hub (bool, default: False) – If True, also upload to HuggingFace Hub

  • **kwargs – Additional arguments for push_to_hub

Return type:

None

Example

>>> model.save_pretrained("./my-tokenizer")
>>> # Later...
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
classmethod from_pretrained(model_name_or_path, map_location=None, **kwargs)[source][source]

Load model from local directory or HuggingFace Hub.

Automatically detects whether path is local or a Hub repository. For Hub repos, downloads config and weights to cache.

Parameters:

  • model_name_or_path (str) – Local path or HuggingFace Hub repo ID

  • map_location (Optional[str], default: None) – Device to load weights to (default: auto-detect)

  • **kwargs – Override config parameters

Return type:

BaseTokenizer

Returns:

Loaded model instance

Example

>>> # From local path
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
>>>
>>> # From HuggingFace Hub
>>> model = ContinuousTokenizer.from_pretrained("username/my-tokenizer")
>>>
>>> # Override config
>>> model = ContinuousTokenizer.from_pretrained(
...     "./my-tokenizer",
...     dropout=0.1  # Override saved dropout value
... )
push_to_hub(repo_id, save_directory=None, commit_message='Upload model', private=False, **kwargs)[source][source]

Upload model to HuggingFace Hub.

Creates or updates a repository on the Hub with model weights and configuration.

Parameters:

  • repo_id (str) – Repository ID (e.g., β€œusername/model-name”)

  • save_directory (Optional[str], default: None) – Local directory to save before upload (auto-created if None)

  • commit_message (str, default: 'Upload model') – Commit message for the upload

  • private (bool, default: False) – Whether to create a private repository

  • **kwargs – Additional arguments for HfApi.upload_folder

Return type:

None

Example

>>> model.push_to_hub("username/my-vae-tokenizer")
>>> # Creates https://huggingface.co/username/my-vae-tokenizer
load_encoder_decoder_weights(pretrained_path, strict=False, verbose=True)[source][source]

Load encoder/decoder weights from a pretrained checkpoint.

This method enables transfer learning by loading encoder and decoder weights from any compatible tokenizer (VAE, VQ-VAE, FSQ, etc.) while leaving quantizer-specific layers randomly initialized.

Use Cases

  1. Initialize discrete tokenizer from VAE: Load MAISI VAE encoder/decoder, then train only the quantizer layers (VQ, FSQ, etc.)

  2. Fine-tune on new domain: Start from pretrained weights, fine-tune all

  3. Encoder-only transfer: Use pretrained encoder for downstream tasks

How It Works

The method identifies encoder/decoder weights by key prefixes and loads only those that match. Quantizer-specific layers (quant_conv, post_quant_conv, quantizer) may or may not be loaded depending on architecture compatibility.

Weight Matching Strategy: - encoder.* keys: Always attempted - decoder.* keys: Always attempted - quant_conv.* keys: Loaded if shapes match - post_quant_conv.* keys: Loaded if shapes match - quantizer.* keys: Skipped (architecture-specific)

type pretrained_path:

str | Path

param pretrained_path:

Path to pretrained weights file (.pt, .bin) or directory containing β€˜pytorch_model.bin’

type strict:

bool, default: False

param strict:

If True, raise error on shape mismatches. If False (default), skip mismatched weights and log warnings.

type verbose:

bool, default: True

param verbose:

If True, print summary of loaded/skipped weights

rtype:

tuple[list[str], list[str]]

returns:

Tuple of (loaded_keys, skipped_keys) for inspection

Example

>>> # Initialize FSQ tokenizer from MAISI VAE weights
>>> model = DiscreteTokenizer(
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     # ... other args matching MAISI architecture
... )
>>> loaded, skipped = model.load_encoder_decoder_weights(
...     'weights/maisi_converted.pt',
...     verbose=True
... )
>>> print(f"Loaded {len(loaded)} keys, skipped {len(skipped)}")
>>>
>>> # Now train with frozen encoder if desired
>>> for param in model.encoder.parameters():
...     param.requires_grad = False

Note

For best results, ensure the pretrained model has the same: - channels, channels_mult, num_res_blocks - spatial_compression

The following can differ: - z_channels, latent_channels, embedding_dim - Quantizer type and configuration

classmethod from_pretrained_encoder_decoder(pretrained_path, strict=False, verbose=True, **kwargs)[source][source]

Create model and load encoder/decoder weights from pretrained checkpoint.

Factory method that creates a new model instance and loads encoder/decoder weights in one step. Useful for initializing new architectures from pretrained backbones.

Parameters:

  • pretrained_path (str | Path) – Path to pretrained weights

  • strict (bool, default: False) – If True, raise on shape mismatches

  • verbose (bool, default: True) – If True, print loading summary

  • **kwargs – Model configuration (passed to __init__)

Return type:

BaseTokenizer

Returns:

New model instance with loaded encoder/decoder weights

Example

>>> # Create FSQ model initialized from VAE weights
>>> model = DiscreteTokenizer.from_pretrained_encoder_decoder(
...     'weights/maisi_converted.pt',
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     z_channels=256,
...     channels=64,
...     channels_mult=(1, 2, 4),
... )
encode_batch(x, batch_size=8, show_progress=False)[source][source]

Encode large batch with automatic mini-batching.

Processes input in chunks to avoid OOM for large datasets.

Parameters:

  • x (Tensor) – Full input tensor of shape (N, C, *spatial)

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated latents for all inputs

Example

>>> dataset = torch.randn(1000, 1, 64, 64, 64)
>>> latents = model.encode_batch(dataset, batch_size=4)
decode_batch(z, batch_size=8, show_progress=False)[source][source]

Decode large batch with automatic mini-batching.

Parameters:

  • z (Tensor) – Full latent tensor

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated reconstructions for all inputs

reconstruct(x, roi_size=None, overlap=0.5, sw_batch_size=1)[source][source]

Full encode-decode reconstruction with optional sliding window.

For large 3D volumes that exceed GPU memory, this method implements sliding window inference with Gaussian importance weighting to seamlessly blend overlapping patches.

The Algorithm

  1. Pad input to multiple of stride + window size

  2. Extract overlapping windows with specified stride

  3. Process windows in batches through tokenize -> detokenize

  4. Weight each window’s contribution by Gaussian importance

  5. Normalize by accumulated importance and crop to original size

Gaussian Weighting

Uses a Gaussian importance map (2D or 3D based on input) that gives higher weight to the center of each window, reducing boundary artifacts when blending.

type x:

Tensor

param x:

Input tensor of shape (B, C, H, W, D) for 3D or (B, C, H, W) for 2D

type roi_size:

Union[tuple[int, ...], int, None], default: None

param roi_size:

Size of sliding window. If None, processes entire volume. Can be int (isotropic) or tuple (anisotropic).

type overlap:

float, default: 0.5

param overlap:

Fraction of overlap between windows (0.0 to 0.9). Higher overlap = smoother blending but more compute.

type sw_batch_size:

int, default: 1

param sw_batch_size:

Number of windows to process in parallel per batch. Higher values use more GPU memory but are faster. Default is 1 (sequential processing).

rtype:

Tensor

returns:

Reconstructed tensor with same shape as input

Example

>>> volume = torch.randn(1, 1, 256, 256, 256)  # 256Β³ volume
>>> # Process in 128Β³ windows with 50% overlap, 4 windows at a time
>>> recon = model.reconstruct(volume, roi_size=128, overlap=0.5, sw_batch_size=4)

Note

  • For volumes that fit in memory, omit roi_size for faster processing

  • Overlap of 0.5 is a good default; higher values reduce artifacts but increase computation proportionally

  • sw_batch_size > 1 can significantly speed up inference on GPUs with sufficient memory

get_latent_shape(input_shape)[source][source]

Calculate latent shape for given input shape.

Parameters:

input_shape (tuple[int, ...]) – Input tensor shape (B, C, *spatial)

Return type:

tuple[int, ...]

Returns:

Expected latent tensor shape

Raises:

NotImplementedError – Must be implemented by subclass

T_destination = ~T_destination
add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

  • name (str) – name of the child module. The child module can be accessed from this module using the given name

  • module (Module) – child module to be added to the module.

Return type:

None

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:

fn (Module -> None) – function to be applied to each submodule

Returns:

self

Return type:

Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor – module buffer

Return type:

Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
call_super_init: bool = False
children()[source]

Return an iterator over immediate children modules.

Yields:

Module – a child module

Return type:

Iterator[Module]

cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

dump_patches: bool = False
eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

self

Return type:

Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The buffer referenced by target

Return type:

torch.Tensor

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Any extra state to store in the module’s state_dict

Return type:

object

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The Parameter referenced by target

Return type:

torch.nn.Parameter

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

Returns:

The submodule referenced by target

Return type:

torch.nn.Module

Raises:

AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

  • state_dict (dict) – a dict containing parameters and persistent buffers.

  • strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

  • assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

  • missing_keys is a list of str containing any keys that are expected

    by this module but missing from the provided state_dict.

  • unexpected_keys is a list of str containing the keys that are not

    expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:

remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.

Yields:

Module – a module in the network

Return type:

Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

  • prefix (str) – prefix to prepend to all buffer names.

  • recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

  • remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Iterator[tuple[str, Tensor]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())
named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Return type:

Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

  • memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result

  • prefix (str, default: '') – a prefix that will be added to the name of the module

  • remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

  • prefix (str) – prefix to prepend to all parameter names.

  • recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

  • remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Iterator[tuple[str, Parameter]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())
parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Return type:

Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

  • name (str) – name of the buffer. The buffer can be accessed from this module using the given name

  • tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.

  • persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

None

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))
register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False

  • always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

  1. Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

  2. If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

  3. If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::

hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:

hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:
Return type:

None

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

  • name (str) – name of the parameter. The parameter can be accessed from this module using the given name

  • param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

None

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::

hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.

Returns:

self

Return type:

Module

set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict

Return type:

None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

  • target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

  • module (Module) – The module to set the submodule to.

  • strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

  • ValueError – If the target string is empty or if module is not an instance of nn.Module.

  • AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

None

share_memory()[source]

See torch.Tensor.share_memory_().

Return type:

Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

  • destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]
to(dtype, non_blocking=False)[source]
to(tensor, non_blocking=False)[source]
to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

  • device (torch.device) – the desired device of the parameters and buffers in this module

  • dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

  • tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

  • memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

  • device (torch.device) – The desired device of the parameters and buffers in this module.

  • recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:

dst_type (type or string) – the desired type

Returns:

self

Return type:

Module

xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:

set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.

Return type:

None

training: bool
class medtokenizers.ContinuousTokenizer(dim, in_channels=1, out_channels=1, z_channels=4, z_factor=None, latent_channels=4, channels=64, channels_mult=(1, 2, 4), num_res_blocks=2, attn_resolutions=(), dropout=0.0, resolution=256, spatial_compression=4, formulation='VAE', use_encoder_mid=False, use_output_nonlinearity=False, decoder_blocks_per_stage=None, separate_quant_conv=True, name='ContinuousTokenizer', **kwargs)[source][source]

Bases: medtokenizers.modules.base.BaseTokenizer

Continuous latent tokenizer for medical imaging (VAE/AE).

This tokenizer learns a continuous latent representation using either: - VAE: Variational Autoencoder with KL divergence regularization - AE: Standard Autoencoder without probabilistic modeling

The VAE variant is particularly useful for: - Latent diffusion models (LDM) - Interpolation in latent space - Generative modeling with controllable sampling

Architecture Details

The network follows a symmetric encoder-decoder design:

Encoder Path:
Input(H,W,D) -> Conv_in -> ResBlocks -> Downsample -> ... -> Conv_out -> mu, sigma^2

Decoder Path:
z -> Conv_in -> ResBlocks -> Upsample -> ... -> Conv_out -> Output(H,W,D)

Key architectural choices: - quant_conv: 1x1 conv reducing encoder output to latent dimension - post_quant_conv: 1x1 conv expanding latent to decoder input - GroupNorm: Batch-size independent normalization - Swish activation: Smooth, non-monotonic activation

Memory Optimization

For 3D volumes, the model automatically: - Uses channels_last_3d memory format for better cache efficiency - Supports gradient checkpointing (via use_checkpointing kwarg)

type dim:

int

param dim:

Spatial dimensionality (2 for 2D images, 3 for 3D volumes)

type in_channels:

int, default: 1

param in_channels:

Number of input channels (1 for grayscale, 3 for RGB)

type out_channels:

int, default: 1

param out_channels:

Number of output channels (usually same as in_channels)

type z_channels:

int, default: 4

param z_channels:

Intermediate channels after encoder, before quant_conv

type z_factor:

Optional[int], default: None

param z_factor:

Multiplier for encoder output channels. Default: 2 for VAE (outputs ΞΌ and σ²), 1 for AE (outputs z directly)

type latent_channels:

int, default: 4

param latent_channels:

Final latent dimension (e.g., 4 for SD-style VAE)

type channels:

int, default: 64

param channels:

Base channel count (scaled by channels_mult)

type channels_mult:

tuple[int, ...], default: (1, 2, 4)

param channels_mult:

Channel multipliers at each resolution level. Example: (1, 2, 4) means channels β†’ 2*channels β†’ 4*channels

type num_res_blocks:

int, default: 2

param num_res_blocks:

Number of residual blocks per resolution level

type attn_resolutions:

tuple[int, ...], default: ()

param attn_resolutions:

Spatial resolutions where self-attention is applied

type dropout:

float, default: 0.0

param dropout:

Dropout probability in residual blocks

type resolution:

int, default: 256

param resolution:

Input spatial resolution (for attention position info)

type spatial_compression:

int, default: 4

param spatial_compression:

Total downsampling factor (e.g., 8 = 3 downsamples)

type formulation:

Literal['VAE', 'AE'], default: 'VAE'

param formulation:

β€œVAE” for variational, β€œAE” for deterministic

type name:

str, default: 'ContinuousTokenizer'

param name:

Model identifier for saving/loading

type **kwargs:

Any

param **kwargs:

Additional args passed to Encoder/Decoder (e.g., use_checkpointing)

Example

>>> # Create a 3D VAE with 4-channel latent (like Stable Diffusion)
>>> model = ContinuousTokenizer(
...     dim=3,
...     in_channels=1,
...     out_channels=1,
...     z_channels=128,
...     latent_channels=4,
...     channels=64,
...     channels_mult=(1, 2, 4),
...     spatial_compression=8,
...     formulation='VAE'
... )
>>>
>>> # Forward pass returns dict with reconstructions and KL loss
>>> volume = torch.randn(1, 1, 128, 128, 128)
>>> output = model(volume)
>>> recon = output['reconstructions']
>>> kl_loss = output.get('kl_loss')  # Only for VAE
>>>
>>> # For inference, use tokenize/detokenize
>>> with model.inference_mode():
...     latents = model.tokenize(volume)  # (1, 4, 16, 16, 16)
...     reconstructed = model.detokenize(latents)

References

Kingma & Welling β€œAuto-Encoding Variational Bayes” (2013) Rombach et al. β€œHigh-Resolution Image Synthesis with Latent Diffusion Models”

param use_encoder_mid:

type use_encoder_mid:

bool, default: False

param use_output_nonlinearity:

type use_output_nonlinearity:

bool, default: False

param decoder_blocks_per_stage:

type decoder_blocks_per_stage:

Optional[list[int]], default: None

param separate_quant_conv:

type separate_quant_conv:

bool, default: True

__init__(dim, in_channels=1, out_channels=1, z_channels=4, z_factor=None, latent_channels=4, channels=64, channels_mult=(1, 2, 4), num_res_blocks=2, attn_resolutions=(), dropout=0.0, resolution=256, spatial_compression=4, formulation='VAE', use_encoder_mid=False, use_output_nonlinearity=False, decoder_blocks_per_stage=None, separate_quant_conv=True, name='ContinuousTokenizer', **kwargs)[source][source]
Parameters:

  • dim (int)

  • in_channels (int, default: 1)

  • out_channels (int, default: 1)

  • z_channels (int, default: 4)

  • z_factor (Optional[int], default: None)

  • latent_channels (int, default: 4)

  • channels (int, default: 64)

  • channels_mult (tuple[int, ...], default: (1, 2, 4))

  • num_res_blocks (int, default: 2)

  • attn_resolutions (tuple[int, ...], default: ())

  • dropout (float, default: 0.0)

  • resolution (int, default: 256)

  • spatial_compression (int, default: 4)

  • formulation (Literal['VAE', 'AE'], default: 'VAE')

  • use_encoder_mid (bool, default: False)

  • use_output_nonlinearity (bool, default: False)

  • decoder_blocks_per_stage (Optional[list[int]], default: None)

  • separate_quant_conv (bool, default: True)

  • name (str, default: 'ContinuousTokenizer')

  • kwargs (Any)

config: dict[str, Any]
encode(x)[source][source]

Encode input to latent representation.

For VAE: - Encoder outputs (ΞΌ, log σ²) - Samples z using reparameterization: z = ΞΌ + Οƒ * Ξ΅ - Returns (z, (kl_loss, (mean, logvar)))

For AE: - Encoder outputs z directly - Returns (z, (zero_kl, zero_logvar))

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input tensor of shape (B, C, *spatial) where: - B: batch size - C: number of channels (must match model’s in_channels) - spatial: (H, W) for 2D or (H, W, D) for 3D

Returns:

  • latent: Sampled or deterministic latent tensor

  • distribution_output: KL loss and posterior parameters

Return type:

Tuple of

Raises:

  • TypeError – If x is not a floating point tensor

  • ValueError – If x has wrong shape or contains NaN/Inf

T_destination = ~T_destination
add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

  • name (str) – name of the child module. The child module can be accessed from this module using the given name

  • module (Module) – child module to be added to the module.

Return type:

None

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:

fn (Module -> None) – function to be applied to each submodule

Returns:

self

Return type:

Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor – module buffer

Return type:

Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
call_super_init: bool = False
children()[source]

Return an iterator over immediate children modules.

Yields:

Module – a child module

Return type:

Iterator[Module]

compile(mode='reduce-overhead', fullgraph=False, **kwargs)[source]

Compile model with torch.compile() for faster inference.

Uses PyTorch 2.0+ compilation to optimize the model graph. The compiled model maintains the same interface but runs faster, especially for repeated inference calls.

Parameters:

  • mode (str, default: 'reduce-overhead') – Compilation mode. Options: - β€œreduce-overhead”: Best for small batches (default) - β€œmax-autotune”: Best throughput, longer warmup - β€œdefault”: Balanced compilation

  • fullgraph (bool, default: False) – If True, require full graph compilation (stricter)

  • **kwargs – Additional arguments passed to torch.compile()

Return type:

BaseTokenizer

Returns:

Compiled model (self, for method chaining)

Example

>>> model = ContinuousTokenizer.from_pretrained("path/to/model")
>>> model = model.compile(mode="reduce-overhead")
>>> # First call triggers compilation (slower)
>>> latents = model.tokenize(batch)
>>> # Subsequent calls are faster
>>> latents = model.tokenize(batch2)

Note

  • Compilation happens on first forward pass

  • Different input shapes may trigger recompilation

  • Use dynamic=False (default) for fixed input sizes

config_name: str = 'config.json'
cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

decode(z)[source][source]

Decode latent representation to output.

Parameters:

z (Float[Tensor, 'batch latent_channels *spatial_compressed']) – Latent tensor from encode() or external source. Shape: (B, latent_channels, *spatial_compressed)

Return type:

Float[Tensor, 'batch channels *spatial']

Returns:

Reconstructed output with original spatial dimensions

Raises:

  • TypeError – If z is not a floating point tensor

  • ValueError – If z has wrong shape or contains NaN/Inf

decode_batch(z, batch_size=8, show_progress=False)[source]

Decode large batch with automatic mini-batching.

Parameters:

  • z (Tensor) – Full latent tensor

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated reconstructions for all inputs

detokenize(z)[source]

Decode from latent space (convenience method).

Parameters:

z (Tensor) – Latent tensor

Return type:

Tensor

Returns:

Reconstructed output

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

dump_patches: bool = False
encode_batch(x, batch_size=8, show_progress=False)[source]

Encode large batch with automatic mini-batching.

Processes input in chunks to avoid OOM for large datasets.

Parameters:

  • x (Tensor) – Full input tensor of shape (N, C, *spatial)

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated latents for all inputs

Example

>>> dataset = torch.randn(1000, 1, 64, 64, 64)
>>> latents = model.encode_batch(dataset, batch_size=4)
eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

self

Return type:

Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

classmethod from_pretrained(model_name_or_path, map_location=None, **kwargs)[source]

Load model from local directory or HuggingFace Hub.

Automatically detects whether path is local or a Hub repository. For Hub repos, downloads config and weights to cache.

Parameters:

  • model_name_or_path (str) – Local path or HuggingFace Hub repo ID

  • map_location (Optional[str], default: None) – Device to load weights to (default: auto-detect)

  • **kwargs – Override config parameters

Return type:

BaseTokenizer

Returns:

Loaded model instance

Example

>>> # From local path
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
>>>
>>> # From HuggingFace Hub
>>> model = ContinuousTokenizer.from_pretrained("username/my-tokenizer")
>>>
>>> # Override config
>>> model = ContinuousTokenizer.from_pretrained(
...     "./my-tokenizer",
...     dropout=0.1  # Override saved dropout value
... )
classmethod from_pretrained_encoder_decoder(pretrained_path, strict=False, verbose=True, **kwargs)[source]

Create model and load encoder/decoder weights from pretrained checkpoint.

Factory method that creates a new model instance and loads encoder/decoder weights in one step. Useful for initializing new architectures from pretrained backbones.

Parameters:

  • pretrained_path (str | Path) – Path to pretrained weights

  • strict (bool, default: False) – If True, raise on shape mismatches

  • verbose (bool, default: True) – If True, print loading summary

  • **kwargs – Model configuration (passed to __init__)

Return type:

BaseTokenizer

Returns:

New model instance with loaded encoder/decoder weights

Example

>>> # Create FSQ model initialized from VAE weights
>>> model = DiscreteTokenizer.from_pretrained_encoder_decoder(
...     'weights/maisi_converted.pt',
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     z_channels=256,
...     channels=64,
...     channels_mult=(1, 2, 4),
... )
get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The buffer referenced by target

Return type:

torch.Tensor

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Any extra state to store in the module’s state_dict

Return type:

object

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The Parameter referenced by target

Return type:

torch.nn.Parameter

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

Returns:

The submodule referenced by target

Return type:

torch.nn.Module

Raises:

AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

inference_mode()[source]

Context manager for optimized inference.

Configures the model for maximum inference performance: - Sets model to eval mode - Disables gradient computation - Uses torch.inference_mode for additional optimizations

The model is restored to its previous state upon exit.

Return type:

Generator[BaseTokenizer, None, None]

Example

>>> with model.inference_mode():
...     latents = model.tokenize(volume)
...     recon = model.detokenize(latents)
Yields:

self – The model instance for method chaining

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

load_encoder_decoder_weights(pretrained_path, strict=False, verbose=True)[source]

Load encoder/decoder weights from a pretrained checkpoint.

This method enables transfer learning by loading encoder and decoder weights from any compatible tokenizer (VAE, VQ-VAE, FSQ, etc.) while leaving quantizer-specific layers randomly initialized.

Use Cases

  1. Initialize discrete tokenizer from VAE: Load MAISI VAE encoder/decoder, then train only the quantizer layers (VQ, FSQ, etc.)

  2. Fine-tune on new domain: Start from pretrained weights, fine-tune all

  3. Encoder-only transfer: Use pretrained encoder for downstream tasks

How It Works

The method identifies encoder/decoder weights by key prefixes and loads only those that match. Quantizer-specific layers (quant_conv, post_quant_conv, quantizer) may or may not be loaded depending on architecture compatibility.

Weight Matching Strategy: - encoder.* keys: Always attempted - decoder.* keys: Always attempted - quant_conv.* keys: Loaded if shapes match - post_quant_conv.* keys: Loaded if shapes match - quantizer.* keys: Skipped (architecture-specific)

type pretrained_path:

str | Path

param pretrained_path:

Path to pretrained weights file (.pt, .bin) or directory containing β€˜pytorch_model.bin’

type strict:

bool, default: False

param strict:

If True, raise error on shape mismatches. If False (default), skip mismatched weights and log warnings.

type verbose:

bool, default: True

param verbose:

If True, print summary of loaded/skipped weights

rtype:

tuple[list[str], list[str]]

returns:

Tuple of (loaded_keys, skipped_keys) for inspection

Example

>>> # Initialize FSQ tokenizer from MAISI VAE weights
>>> model = DiscreteTokenizer(
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     # ... other args matching MAISI architecture
... )
>>> loaded, skipped = model.load_encoder_decoder_weights(
...     'weights/maisi_converted.pt',
...     verbose=True
... )
>>> print(f"Loaded {len(loaded)} keys, skipped {len(skipped)}")
>>>
>>> # Now train with frozen encoder if desired
>>> for param in model.encoder.parameters():
...     param.requires_grad = False

Note

For best results, ensure the pretrained model has the same: - channels, channels_mult, num_res_blocks - spatial_compression

The following can differ: - z_channels, latent_channels, embedding_dim - Quantizer type and configuration

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

  • state_dict (dict) – a dict containing parameters and persistent buffers.

  • strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

  • assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

  • missing_keys is a list of str containing any keys that are expected

    by this module but missing from the provided state_dict.

  • unexpected_keys is a list of str containing the keys that are not

    expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:

remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.

Yields:

Module – a module in the network

Return type:

Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

  • prefix (str) – prefix to prepend to all buffer names.

  • recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

  • remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Iterator[tuple[str, Tensor]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())
named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Return type:

Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

  • memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result

  • prefix (str, default: '') – a prefix that will be added to the name of the module

  • remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

  • prefix (str) – prefix to prepend to all parameter names.

  • recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

  • remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Iterator[tuple[str, Parameter]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())
num_parameters()[source]

Get total number of learnable parameters.

Return type:

int

Returns:

Sum of numel() for all parameters

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Return type:

Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
push_to_hub(repo_id, save_directory=None, commit_message='Upload model', private=False, **kwargs)[source]

Upload model to HuggingFace Hub.

Creates or updates a repository on the Hub with model weights and configuration.

Parameters:

  • repo_id (str) – Repository ID (e.g., β€œusername/model-name”)

  • save_directory (Optional[str], default: None) – Local directory to save before upload (auto-created if None)

  • commit_message (str, default: 'Upload model') – Commit message for the upload

  • private (bool, default: False) – Whether to create a private repository

  • **kwargs – Additional arguments for HfApi.upload_folder

Return type:

None

Example

>>> model.push_to_hub("username/my-vae-tokenizer")
>>> # Creates https://huggingface.co/username/my-vae-tokenizer
reconstruct(x, roi_size=None, overlap=0.5, sw_batch_size=1)[source]

Full encode-decode reconstruction with optional sliding window.

For large 3D volumes that exceed GPU memory, this method implements sliding window inference with Gaussian importance weighting to seamlessly blend overlapping patches.

The Algorithm

  1. Pad input to multiple of stride + window size

  2. Extract overlapping windows with specified stride

  3. Process windows in batches through tokenize -> detokenize

  4. Weight each window’s contribution by Gaussian importance

  5. Normalize by accumulated importance and crop to original size

Gaussian Weighting

Uses a Gaussian importance map (2D or 3D based on input) that gives higher weight to the center of each window, reducing boundary artifacts when blending.

type x:

Tensor

param x:

Input tensor of shape (B, C, H, W, D) for 3D or (B, C, H, W) for 2D

type roi_size:

Union[tuple[int, ...], int, None], default: None

param roi_size:

Size of sliding window. If None, processes entire volume. Can be int (isotropic) or tuple (anisotropic).

type overlap:

float, default: 0.5

param overlap:

Fraction of overlap between windows (0.0 to 0.9). Higher overlap = smoother blending but more compute.

type sw_batch_size:

int, default: 1

param sw_batch_size:

Number of windows to process in parallel per batch. Higher values use more GPU memory but are faster. Default is 1 (sequential processing).

rtype:

Tensor

returns:

Reconstructed tensor with same shape as input

Example

>>> volume = torch.randn(1, 1, 256, 256, 256)  # 256Β³ volume
>>> # Process in 128Β³ windows with 50% overlap, 4 windows at a time
>>> recon = model.reconstruct(volume, roi_size=128, overlap=0.5, sw_batch_size=4)

Note

  • For volumes that fit in memory, omit roi_size for faster processing

  • Overlap of 0.5 is a good default; higher values reduce artifacts but increase computation proportionally

  • sw_batch_size > 1 can significantly speed up inference on GPUs with sufficient memory

register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

  • name (str) – name of the buffer. The buffer can be accessed from this module using the given name

  • tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.

  • persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

None

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))
register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False

  • always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

  1. Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

  2. If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

  3. If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::

hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:

hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:
Return type:

None

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

  • name (str) – name of the parameter. The parameter can be accessed from this module using the given name

  • param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

None

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::

hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.

Returns:

self

Return type:

Module

save_pretrained(save_directory, push_to_hub=False, **kwargs)[source]

Save model weights and configuration to directory.

Creates a directory structure compatible with from_pretrained(): ` save_directory/ β”œβ”€β”€ config.json      # Model configuration └── pytorch_model.bin # Model weights `

Parameters:

  • save_directory (str | Path) – Path to save model

  • push_to_hub (bool, default: False) – If True, also upload to HuggingFace Hub

  • **kwargs – Additional arguments for push_to_hub

Return type:

None

Example

>>> model.save_pretrained("./my-tokenizer")
>>> # Later...
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict

Return type:

None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

  • target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

  • module (Module) – The module to set the submodule to.

  • strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

  • ValueError – If the target string is empty or if module is not an instance of nn.Module.

  • AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

None

share_memory()[source]

See torch.Tensor.share_memory_().

Return type:

Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

  • destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]
to(dtype, non_blocking=False)[source]
to(tensor, non_blocking=False)[source]
to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

  • device (torch.device) – the desired device of the parameters and buffers in this module

  • dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

  • tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

  • memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

  • device (torch.device) – The desired device of the parameters and buffers in this module.

  • recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:

dst_type (type or string) – the desired type

Returns:

self

Return type:

Module

weights_name: str = 'pytorch_model.bin'
xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:

set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.

Return type:

None

training: bool
forward(input)[source][source]

Full forward pass: encode -> decode.

During training, returns a dict with all outputs for loss computation. During evaluation, returns a NetworkEval namedtuple.

Parameters:

input (Float[Tensor, 'batch channels *spatial']) – Input tensor

Returns:

  • β€˜reconstructions’: Decoded output

    • ’posteriors’: (mean, logvar) for VAE, or identity for AE

    • ’latent’/’latents’: Sampled latent tensor

    • ’kl_loss’: KL divergence (VAE only)

Eval mode (NetworkEval):

  • reconstructions: Decoded output

  • posteriors: Distribution parameters

  • latent: Sampled latent

Return type:

Training mode (dict)

tokenize(x)[source][source]

Encode input to latent representation (convenience method).

For inference, this is the primary encoding method. For VAE, returns the sampled latent (not the mean), enabling diverse reconstructions.

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input tensor

Return type:

Float[Tensor, 'batch latent_channels *spatial_compressed']

Returns:

Latent tensor suitable for storage, manipulation, or decoding

get_latent_shape(input_shape)[source][source]

Calculate output latent shape for given input shape.

Useful for pre-allocating memory or understanding compression ratio.

Parameters:

input_shape (tuple[int, ...]) – Input tensor shape (B, C, H, W) or (B, C, H, W, D)

Return type:

tuple[int, ...]

Returns:

Expected latent shape (B, latent_channels, H’, W’) or (B, C’, H’, W’, D’) where spatial dims are compressed by spatial_compression factor

class medtokenizers.DiscreteTokenizer(dim, in_channels=1, out_channels=1, z_channels=4, embedding_dim=6, channels=64, channels_mult=(1, 2, 4), num_res_blocks=2, attn_resolutions=(), dropout=0.0, resolution=256, spatial_compression=4, quantizer='RESFSQ', use_encoder_mid=False, use_output_nonlinearity=False, decoder_blocks_per_stage=None, num_embeddings=1024, beta=0.25, use_norm=False, use_ema=False, ema_decay=0.99, levels=None, num_codebooks=1, codebook_size=None, codebook_dim=None, entropy_loss_weight=0.1, commitment_loss_weight=0.25, quant_temp=0.01, name='DiscreteTokenizer', **kwargs)[source][source]

Bases: medtokenizers.modules.base.BaseTokenizer

Discrete latent tokenizer for medical imaging.

This tokenizer learns discrete latent representations using various quantization methods, enabling the use of language models and autoregressive architectures for medical image generation.

The Key Insight

By quantizing continuous encoder outputs to a finite vocabulary, we convert the image generation problem into a sequence modeling problem that can leverage powerful transformer architectures.

Quantization Methods

Choose based on your use case:

  • VQ: Maximum expressiveness, but requires careful training to avoid codebook collapse. Best for small codebooks (~1K).

  • FSQ: Stable training with implicit codebook. No collapse. Good default choice for most applications.

  • RESFSQ: Massive effective codebook via residual stacking. Use when you need very high fidelity reconstruction.

  • LFQ: Binary codes for extreme simplicity. Good for very large-scale generation with lightweight decoders.

Architecture Details

Encoder Path:
Input -> Conv_in -> ResBlocks -> Downsample -> Conv_out -> z_continuous

Quantization:
z_continuous -> quant_conv -> Quantizer -> (indices, z_quantized)

Decoder Path:
z_quantized -> post_quant_conv -> ResBlocks -> Upsample -> Output

Memory Optimization

For 3D volumes, the model automatically: - Uses channels_last_3d memory format - Supports gradient checkpointing

type dim:

int

param dim:

Spatial dimensionality (2 for 2D, 3 for 3D)

type in_channels:

int, default: 1

param in_channels:

Number of input channels

type out_channels:

int, default: 1

param out_channels:

Number of output channels

type z_channels:

int, default: 4

param z_channels:

Encoder output channels (before quant_conv)

type embedding_dim:

int, default: 6

param embedding_dim:

Dimension of quantized embeddings

type channels:

int, default: 64

param channels:

Base channel count for encoder/decoder

type channels_mult:

tuple[int, ...], default: (1, 2, 4)

param channels_mult:

Channel multipliers per resolution

type num_res_blocks:

int, default: 2

param num_res_blocks:

Residual blocks per resolution

type attn_resolutions:

tuple[int, ...], default: ()

param attn_resolutions:

Resolutions for self-attention

type dropout:

float, default: 0.0

param dropout:

Dropout probability

type resolution:

int, default: 256

param resolution:

Input spatial resolution

type spatial_compression:

int, default: 4

param spatial_compression:

Total downsampling factor

type quantizer:

Literal['VQ', 'FSQ', 'LFQ', 'RESFSQ'], default: 'RESFSQ'

param quantizer:

Quantization method (β€œVQ”, β€œFSQ”, β€œLFQ”, β€œRESFSQ”)

type num_embeddings:

int, default: 1024

param num_embeddings:

Codebook size for VQ (default: 1024)

type beta:

float, default: 0.25

param beta:

Commitment loss weight for VQ (default: 0.25)

type use_norm:

bool, default: False

param use_norm:

Normalize VQ embeddings (cosine similarity)

type levels:

list[int] | None, default: None

param levels:

FSQ quantization levels (e.g., [8, 5, 5, 5])

type num_codebooks:

int, default: 1

param num_codebooks:

Number of quantizers for RESFSQ/LFQ

type codebook_size:

Optional[int], default: None

param codebook_size:

LFQ codebook size (must be power of 2)

type codebook_dim:

Optional[int], default: None

param codebook_dim:

LFQ code dimension

type entropy_loss_weight:

float, default: 0.1

param entropy_loss_weight:

LFQ entropy regularization weight

type commitment_loss_weight:

float, default: 0.25

param commitment_loss_weight:

LFQ commitment loss weight

type quant_temp:

float, default: 0.01

param quant_temp:

Temperature for soft quantization

type name:

str, default: 'DiscreteTokenizer'

param name:

Model identifier

type **kwargs:

Any

param **kwargs:

Additional encoder/decoder arguments

Example

>>> # FSQ tokenizer for 3D medical volumes
>>> model = DiscreteTokenizer(
...     dim=3,
...     in_channels=1,
...     out_channels=1,
...     z_channels=128,
...     embedding_dim=6,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],  # 1000 codes
...     spatial_compression=8,
... )
>>>
>>> # Tokenize to discrete codes
>>> volume = torch.randn(1, 1, 128, 128, 128)
>>> with model.inference_mode():
...     indices = model.tokenize(volume)  # (1, 16, 16, 16)
...     reconstructed = model.detokenize(indices)
>>>
>>> # Training forward pass
>>> output = model(volume)
>>> recon_loss = F.l1_loss(output['reconstructions'], volume)
>>> quant_loss = output['quant_loss'].mean()
>>> total_loss = recon_loss + quant_loss

References

van den Oord et al. β€œNeural Discrete Representation Learning” (VQ-VAE) Mentzer et al. β€œFinite Scalar Quantization: VQ-VAE Made Simple” (FSQ) Yu et al. β€œLanguage Model Beats Diffusion” (LFQ in MagViT-2)

param use_encoder_mid:

type use_encoder_mid:

bool, default: False

param use_output_nonlinearity:

type use_output_nonlinearity:

bool, default: False

param decoder_blocks_per_stage:

type decoder_blocks_per_stage:

Optional[list[int]], default: None

param use_ema:

type use_ema:

bool, default: False

param ema_decay:

type ema_decay:

float, default: 0.99

__init__(dim, in_channels=1, out_channels=1, z_channels=4, embedding_dim=6, channels=64, channels_mult=(1, 2, 4), num_res_blocks=2, attn_resolutions=(), dropout=0.0, resolution=256, spatial_compression=4, quantizer='RESFSQ', use_encoder_mid=False, use_output_nonlinearity=False, decoder_blocks_per_stage=None, num_embeddings=1024, beta=0.25, use_norm=False, use_ema=False, ema_decay=0.99, levels=None, num_codebooks=1, codebook_size=None, codebook_dim=None, entropy_loss_weight=0.1, commitment_loss_weight=0.25, quant_temp=0.01, name='DiscreteTokenizer', **kwargs)[source][source]
Parameters:

  • dim (int)

  • in_channels (int, default: 1)

  • out_channels (int, default: 1)

  • z_channels (int, default: 4)

  • embedding_dim (int, default: 6)

  • channels (int, default: 64)

  • channels_mult (tuple[int, ...], default: (1, 2, 4))

  • num_res_blocks (int, default: 2)

  • attn_resolutions (tuple[int, ...], default: ())

  • dropout (float, default: 0.0)

  • resolution (int, default: 256)

  • spatial_compression (int, default: 4)

  • quantizer (Literal['VQ', 'FSQ', 'LFQ', 'RESFSQ'], default: 'RESFSQ')

  • use_encoder_mid (bool, default: False)

  • use_output_nonlinearity (bool, default: False)

  • decoder_blocks_per_stage (Optional[list[int]], default: None)

  • num_embeddings (int, default: 1024)

  • beta (float, default: 0.25)

  • use_norm (bool, default: False)

  • use_ema (bool, default: False)

  • ema_decay (float, default: 0.99)

  • levels (list[int] | None, default: None)

  • num_codebooks (int, default: 1)

  • codebook_size (Optional[int], default: None)

  • codebook_dim (Optional[int], default: None)

  • entropy_loss_weight (float, default: 0.1)

  • commitment_loss_weight (float, default: 0.25)

  • quant_temp (float, default: 0.01)

  • name (str, default: 'DiscreteTokenizer')

  • kwargs (Any)

config: dict[str, Any]
to(*args, **kwargs)[source][source]

Move and/or cast the model, keeping the quantizer dtype in sync.

The quantizer keeps its own dtype attribute (used by its numerical guards). It is updated only when a dtype is actually supplied, so a plain device move such as model.to("cuda") no longer silently resets it to float32. A dtype may be passed either positionally (model.to(torch.float16)) or as the dtype keyword.

Parameters:
Return type:

DiscreteTokenizer

Returns:

self, after the move/cast has been applied.

encode(x)[source][source]

Encode input to discrete codes.

Passes input through encoder, projects to embedding dimension, then quantizes to discrete codebook indices.

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input tensor of shape (B, C, *spatial) where: - B: batch size - C: number of channels (must match model’s in_channels) - spatial: (H, W) for 2D or (H, W, D) for 3D

Returns:

  • indices: Discrete codebook indices

  • quantized: Quantized continuous codes (for decoder)

  • loss: Quantization loss (commitment, entropy, etc.)

Return type:

Tuple of

Raises:

  • TypeError – If x is not a floating point tensor

  • ValueError – If x has wrong shape or contains NaN/Inf

decode(quant)[source][source]

Decode from quantized continuous codes.

Parameters:

quant (Float[Tensor, 'batch embedding_dim *spatial_compressed']) – Quantized codes from encode() (continuous representation). Shape: (B, embedding_dim, *spatial_compressed)

Return type:

Float[Tensor, 'batch channels *spatial']

Returns:

Reconstructed output with original spatial dimensions

Raises:

  • TypeError – If quant is not a floating point tensor

  • ValueError – If quant has wrong shape or contains NaN/Inf

forward(input)[source][source]

Full forward pass: encode -> quantize -> decode.

During training, returns dict with all outputs for loss computation. During evaluation, returns NetworkEval namedtuple.

Parameters:

input (Float[Tensor, 'batch channels *spatial']) – Input tensor

Returns:

  • β€˜reconstructions’: Decoded output

    • ’quant_loss’: Quantization loss

    • ’quant_info’: Discrete indices

    • ’latents’: Quantized codes (continuous)

Eval mode (NetworkEval):

  • reconstructions: Decoded output

  • quant_loss: Quantization loss

  • quant_info: Discrete indices

Return type:

Training mode (dict)

tokenize(x)[source][source]

Encode input to discrete token indices.

This is the primary encoding method for inference and storage. Returns integer indices that can be stored efficiently or fed to autoregressive models.

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input tensor

Return type:

Int[Tensor, 'batch *spatial_indices']

Returns:

Discrete indices suitable for storage or sequence modeling

detokenize(indices, spatial_shape=None)[source][source]

Decode from discrete token indices.

Inverse of tokenize(). Converts discrete indices back to continuous output via codebook lookup and decoder. This is the canonical index-to-reconstruction decoding path for inference, converting stored/generated indices back to images.

Parameters:

  • indices (Int[Tensor, 'batch *spatial_indices']) – Discrete indices from tokenize(). Can be: - Spatial format: (B, H’, W’) for 2D or (B, H’, W’, D’) for 3D - Flattened format: (B, N) where N = H’ * W’ [* D’]

  • spatial_shape (tuple[int, ...] | None, default: None) – Original latent spatial dimensions (H’, W’) or (H’, W’, D’). Required when indices are flattened to avoid incorrect cubic assumptions for anisotropic volumes.

Return type:

Float[Tensor, 'batch channels *spatial']

Returns:

Reconstructed output

Raises:

ValueError – If spatial_shape is required but not provided

get_latent_shape(input_shape)[source][source]

Calculate discrete index shape for given input shape.

Note: Returns shape WITHOUT embedding dimension (just spatial). For RESFSQ, includes num_quantizers dimension.

Parameters:

input_shape (tuple[int, ...]) – Input tensor shape (B, C, H, W) or (B, C, H, W, D)

Return type:

tuple[int, ...]

Returns:

Expected index shape (B, H’, W’) or (B, H’, W’, D’) where spatial dims are compressed by spatial_compression

get_codebook_size()[source][source]

Get the total vocabulary size.

Return type:

int

Returns:

Number of discrete codes in vocabulary

reconstruct(x, roi_size=None, overlap=0.0)[source][source]

Reconstruction with optional sliding window (overlap must be 0.0).

Discrete tokenizers don’t support overlapping windows because averaging discrete codes is not meaningful.

Parameters:

  • x (Tensor) – Input tensor

  • roi_size (Union[tuple[int, ...], int, None], default: None) – Window size for sliding window inference

  • overlap (float, default: 0.0) – Must be 0.0 for discrete tokenizers

Return type:

Tensor

Returns:

Reconstructed output

Raises:

ValueError – If overlap > 0.0

T_destination = ~T_destination
add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

  • name (str) – name of the child module. The child module can be accessed from this module using the given name

  • module (Module) – child module to be added to the module.

Return type:

None

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:

fn (Module -> None) – function to be applied to each submodule

Returns:

self

Return type:

Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor – module buffer

Return type:

Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
call_super_init: bool = False
children()[source]

Return an iterator over immediate children modules.

Yields:

Module – a child module

Return type:

Iterator[Module]

compile(mode='reduce-overhead', fullgraph=False, **kwargs)[source]

Compile model with torch.compile() for faster inference.

Uses PyTorch 2.0+ compilation to optimize the model graph. The compiled model maintains the same interface but runs faster, especially for repeated inference calls.

Parameters:

  • mode (str, default: 'reduce-overhead') – Compilation mode. Options: - β€œreduce-overhead”: Best for small batches (default) - β€œmax-autotune”: Best throughput, longer warmup - β€œdefault”: Balanced compilation

  • fullgraph (bool, default: False) – If True, require full graph compilation (stricter)

  • **kwargs – Additional arguments passed to torch.compile()

Return type:

BaseTokenizer

Returns:

Compiled model (self, for method chaining)

Example

>>> model = ContinuousTokenizer.from_pretrained("path/to/model")
>>> model = model.compile(mode="reduce-overhead")
>>> # First call triggers compilation (slower)
>>> latents = model.tokenize(batch)
>>> # Subsequent calls are faster
>>> latents = model.tokenize(batch2)

Note

  • Compilation happens on first forward pass

  • Different input shapes may trigger recompilation

  • Use dynamic=False (default) for fixed input sizes

config_name: str = 'config.json'
cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

decode_batch(z, batch_size=8, show_progress=False)[source]

Decode large batch with automatic mini-batching.

Parameters:

  • z (Tensor) – Full latent tensor

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated reconstructions for all inputs

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

dump_patches: bool = False
encode_batch(x, batch_size=8, show_progress=False)[source]

Encode large batch with automatic mini-batching.

Processes input in chunks to avoid OOM for large datasets.

Parameters:

  • x (Tensor) – Full input tensor of shape (N, C, *spatial)

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated latents for all inputs

Example

>>> dataset = torch.randn(1000, 1, 64, 64, 64)
>>> latents = model.encode_batch(dataset, batch_size=4)
eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

self

Return type:

Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

classmethod from_pretrained(model_name_or_path, map_location=None, **kwargs)[source]

Load model from local directory or HuggingFace Hub.

Automatically detects whether path is local or a Hub repository. For Hub repos, downloads config and weights to cache.

Parameters:

  • model_name_or_path (str) – Local path or HuggingFace Hub repo ID

  • map_location (Optional[str], default: None) – Device to load weights to (default: auto-detect)

  • **kwargs – Override config parameters

Return type:

BaseTokenizer

Returns:

Loaded model instance

Example

>>> # From local path
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
>>>
>>> # From HuggingFace Hub
>>> model = ContinuousTokenizer.from_pretrained("username/my-tokenizer")
>>>
>>> # Override config
>>> model = ContinuousTokenizer.from_pretrained(
...     "./my-tokenizer",
...     dropout=0.1  # Override saved dropout value
... )
classmethod from_pretrained_encoder_decoder(pretrained_path, strict=False, verbose=True, **kwargs)[source]

Create model and load encoder/decoder weights from pretrained checkpoint.

Factory method that creates a new model instance and loads encoder/decoder weights in one step. Useful for initializing new architectures from pretrained backbones.

Parameters:

  • pretrained_path (str | Path) – Path to pretrained weights

  • strict (bool, default: False) – If True, raise on shape mismatches

  • verbose (bool, default: True) – If True, print loading summary

  • **kwargs – Model configuration (passed to __init__)

Return type:

BaseTokenizer

Returns:

New model instance with loaded encoder/decoder weights

Example

>>> # Create FSQ model initialized from VAE weights
>>> model = DiscreteTokenizer.from_pretrained_encoder_decoder(
...     'weights/maisi_converted.pt',
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     z_channels=256,
...     channels=64,
...     channels_mult=(1, 2, 4),
... )
get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The buffer referenced by target

Return type:

torch.Tensor

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Any extra state to store in the module’s state_dict

Return type:

object

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The Parameter referenced by target

Return type:

torch.nn.Parameter

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

Returns:

The submodule referenced by target

Return type:

torch.nn.Module

Raises:

AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

inference_mode()[source]

Context manager for optimized inference.

Configures the model for maximum inference performance: - Sets model to eval mode - Disables gradient computation - Uses torch.inference_mode for additional optimizations

The model is restored to its previous state upon exit.

Return type:

Generator[BaseTokenizer, None, None]

Example

>>> with model.inference_mode():
...     latents = model.tokenize(volume)
...     recon = model.detokenize(latents)
Yields:

self – The model instance for method chaining

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

load_encoder_decoder_weights(pretrained_path, strict=False, verbose=True)[source]

Load encoder/decoder weights from a pretrained checkpoint.

This method enables transfer learning by loading encoder and decoder weights from any compatible tokenizer (VAE, VQ-VAE, FSQ, etc.) while leaving quantizer-specific layers randomly initialized.

Use Cases

  1. Initialize discrete tokenizer from VAE: Load MAISI VAE encoder/decoder, then train only the quantizer layers (VQ, FSQ, etc.)

  2. Fine-tune on new domain: Start from pretrained weights, fine-tune all

  3. Encoder-only transfer: Use pretrained encoder for downstream tasks

How It Works

The method identifies encoder/decoder weights by key prefixes and loads only those that match. Quantizer-specific layers (quant_conv, post_quant_conv, quantizer) may or may not be loaded depending on architecture compatibility.

Weight Matching Strategy: - encoder.* keys: Always attempted - decoder.* keys: Always attempted - quant_conv.* keys: Loaded if shapes match - post_quant_conv.* keys: Loaded if shapes match - quantizer.* keys: Skipped (architecture-specific)

type pretrained_path:

str | Path

param pretrained_path:

Path to pretrained weights file (.pt, .bin) or directory containing β€˜pytorch_model.bin’

type strict:

bool, default: False

param strict:

If True, raise error on shape mismatches. If False (default), skip mismatched weights and log warnings.

type verbose:

bool, default: True

param verbose:

If True, print summary of loaded/skipped weights

rtype:

tuple[list[str], list[str]]

returns:

Tuple of (loaded_keys, skipped_keys) for inspection

Example

>>> # Initialize FSQ tokenizer from MAISI VAE weights
>>> model = DiscreteTokenizer(
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     # ... other args matching MAISI architecture
... )
>>> loaded, skipped = model.load_encoder_decoder_weights(
...     'weights/maisi_converted.pt',
...     verbose=True
... )
>>> print(f"Loaded {len(loaded)} keys, skipped {len(skipped)}")
>>>
>>> # Now train with frozen encoder if desired
>>> for param in model.encoder.parameters():
...     param.requires_grad = False

Note

For best results, ensure the pretrained model has the same: - channels, channels_mult, num_res_blocks - spatial_compression

The following can differ: - z_channels, latent_channels, embedding_dim - Quantizer type and configuration

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

  • state_dict (dict) – a dict containing parameters and persistent buffers.

  • strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

  • assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

  • missing_keys is a list of str containing any keys that are expected

    by this module but missing from the provided state_dict.

  • unexpected_keys is a list of str containing the keys that are not

    expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:

remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.

Yields:

Module – a module in the network

Return type:

Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

  • prefix (str) – prefix to prepend to all buffer names.

  • recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

  • remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Iterator[tuple[str, Tensor]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())
named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Return type:

Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

  • memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result

  • prefix (str, default: '') – a prefix that will be added to the name of the module

  • remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

  • prefix (str) – prefix to prepend to all parameter names.

  • recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

  • remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Iterator[tuple[str, Parameter]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())
num_parameters()[source]

Get total number of learnable parameters.

Return type:

int

Returns:

Sum of numel() for all parameters

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Return type:

Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
push_to_hub(repo_id, save_directory=None, commit_message='Upload model', private=False, **kwargs)[source]

Upload model to HuggingFace Hub.

Creates or updates a repository on the Hub with model weights and configuration.

Parameters:

  • repo_id (str) – Repository ID (e.g., β€œusername/model-name”)

  • save_directory (Optional[str], default: None) – Local directory to save before upload (auto-created if None)

  • commit_message (str, default: 'Upload model') – Commit message for the upload

  • private (bool, default: False) – Whether to create a private repository

  • **kwargs – Additional arguments for HfApi.upload_folder

Return type:

None

Example

>>> model.push_to_hub("username/my-vae-tokenizer")
>>> # Creates https://huggingface.co/username/my-vae-tokenizer
register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

  • name (str) – name of the buffer. The buffer can be accessed from this module using the given name

  • tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.

  • persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

None

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))
register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False

  • always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

  1. Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

  2. If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

  3. If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::

hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:

hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:
Return type:

None

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

  • name (str) – name of the parameter. The parameter can be accessed from this module using the given name

  • param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

None

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::

hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.

Returns:

self

Return type:

Module

save_pretrained(save_directory, push_to_hub=False, **kwargs)[source]

Save model weights and configuration to directory.

Creates a directory structure compatible with from_pretrained(): ` save_directory/ β”œβ”€β”€ config.json      # Model configuration └── pytorch_model.bin # Model weights `

Parameters:

  • save_directory (str | Path) – Path to save model

  • push_to_hub (bool, default: False) – If True, also upload to HuggingFace Hub

  • **kwargs – Additional arguments for push_to_hub

Return type:

None

Example

>>> model.save_pretrained("./my-tokenizer")
>>> # Later...
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict

Return type:

None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

  • target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

  • module (Module) – The module to set the submodule to.

  • strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

  • ValueError – If the target string is empty or if module is not an instance of nn.Module.

  • AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

None

share_memory()[source]

See torch.Tensor.share_memory_().

Return type:

Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

  • destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

  • device (torch.device) – The desired device of the parameters and buffers in this module.

  • recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:

dst_type (type or string) – the desired type

Returns:

self

Return type:

Module

weights_name: str = 'pytorch_model.bin'
xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:

set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.

Return type:

None

training: bool
class medtokenizers.MAISITokenizer(pretrained=None, **kwargs)[source][source]

Bases: medtokenizers.networks.continuous.ContinuousTokenizer

MAISI VAE tokenizer matching NVIDIA NV-Generate-MR architecture.

To use published NVIDIA MAISI weights, first convert them with scripts/convert_maisi_to_hf.py and then load the converted checkpoint via from_pretrained(). Note that NVIDIA MAISI model weights are distributed under the NVIDIA Source Code License (NSCLv1), separate from this repository’s Apache-2.0 code license.

Parameters:

pretrained (Optional[str], default: None)

DEFAULT_CONFIG = {'attn_resolutions': (), 'channels': 64, 'channels_mult': (1, 2, 4), 'decoder_blocks_per_stage': [2, 2, 0], 'dim': 3, 'dropout': 0.0, 'formulation': 'VAE', 'in_channels': 1, 'latent_channels': 4, 'name': 'MAISITokenizer', 'num_res_blocks': 2, 'out_channels': 1, 'resolution': 256, 'spatial_compression': 4, 'use_encoder_mid': False, 'z_channels': 4, 'z_factor': 2}
__init__(pretrained=None, **kwargs)[source][source]
Parameters:

pretrained (Optional[str], default: None)

classmethod from_pretrained(pretrained_model_name_or_path, **kwargs)[source][source]

Load from local checkpoint.

Parameters:

pretrained_model_name_or_path (str)

Return type:

MAISITokenizer

static get_training_config()[source][source]

Recommended training hyperparameters from MAISI paper.

Return type:

dict

T_destination = ~T_destination
add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

  • name (str) – name of the child module. The child module can be accessed from this module using the given name

  • module (Module) – child module to be added to the module.

Return type:

None

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:

fn (Module -> None) – function to be applied to each submodule

Returns:

self

Return type:

Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor – module buffer

Return type:

Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
call_super_init: bool = False
children()[source]

Return an iterator over immediate children modules.

Yields:

Module – a child module

Return type:

Iterator[Module]

compile(mode='reduce-overhead', fullgraph=False, **kwargs)[source]

Compile model with torch.compile() for faster inference.

Uses PyTorch 2.0+ compilation to optimize the model graph. The compiled model maintains the same interface but runs faster, especially for repeated inference calls.

Parameters:

  • mode (str, default: 'reduce-overhead') – Compilation mode. Options: - β€œreduce-overhead”: Best for small batches (default) - β€œmax-autotune”: Best throughput, longer warmup - β€œdefault”: Balanced compilation

  • fullgraph (bool, default: False) – If True, require full graph compilation (stricter)

  • **kwargs – Additional arguments passed to torch.compile()

Return type:

BaseTokenizer

Returns:

Compiled model (self, for method chaining)

Example

>>> model = ContinuousTokenizer.from_pretrained("path/to/model")
>>> model = model.compile(mode="reduce-overhead")
>>> # First call triggers compilation (slower)
>>> latents = model.tokenize(batch)
>>> # Subsequent calls are faster
>>> latents = model.tokenize(batch2)

Note

  • Compilation happens on first forward pass

  • Different input shapes may trigger recompilation

  • Use dynamic=False (default) for fixed input sizes

config_name: str = 'config.json'
cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

decode(z)[source]

Decode latent representation to output.

Parameters:

z (Float[Tensor, 'batch latent_channels *spatial_compressed']) – Latent tensor from encode() or external source. Shape: (B, latent_channels, *spatial_compressed)

Return type:

Float[Tensor, 'batch channels *spatial']

Returns:

Reconstructed output with original spatial dimensions

Raises:

  • TypeError – If z is not a floating point tensor

  • ValueError – If z has wrong shape or contains NaN/Inf

decode_batch(z, batch_size=8, show_progress=False)[source]

Decode large batch with automatic mini-batching.

Parameters:

  • z (Tensor) – Full latent tensor

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated reconstructions for all inputs

detokenize(z)[source]

Decode from latent space (convenience method).

Parameters:

z (Tensor) – Latent tensor

Return type:

Tensor

Returns:

Reconstructed output

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

dump_patches: bool = False
encode(x)[source]

Encode input to latent representation.

For VAE: - Encoder outputs (ΞΌ, log σ²) - Samples z using reparameterization: z = ΞΌ + Οƒ * Ξ΅ - Returns (z, (kl_loss, (mean, logvar)))

For AE: - Encoder outputs z directly - Returns (z, (zero_kl, zero_logvar))

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input tensor of shape (B, C, *spatial) where: - B: batch size - C: number of channels (must match model’s in_channels) - spatial: (H, W) for 2D or (H, W, D) for 3D

Returns:

  • latent: Sampled or deterministic latent tensor

  • distribution_output: KL loss and posterior parameters

Return type:

Tuple of

Raises:

  • TypeError – If x is not a floating point tensor

  • ValueError – If x has wrong shape or contains NaN/Inf

encode_batch(x, batch_size=8, show_progress=False)[source]

Encode large batch with automatic mini-batching.

Processes input in chunks to avoid OOM for large datasets.

Parameters:

  • x (Tensor) – Full input tensor of shape (N, C, *spatial)

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated latents for all inputs

Example

>>> dataset = torch.randn(1000, 1, 64, 64, 64)
>>> latents = model.encode_batch(dataset, batch_size=4)
eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

self

Return type:

Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

forward(input)[source]

Full forward pass: encode -> decode.

During training, returns a dict with all outputs for loss computation. During evaluation, returns a NetworkEval namedtuple.

Parameters:

input (Float[Tensor, 'batch channels *spatial']) – Input tensor

Returns:

  • β€˜reconstructions’: Decoded output

    • ’posteriors’: (mean, logvar) for VAE, or identity for AE

    • ’latent’/’latents’: Sampled latent tensor

    • ’kl_loss’: KL divergence (VAE only)

Eval mode (NetworkEval):

  • reconstructions: Decoded output

  • posteriors: Distribution parameters

  • latent: Sampled latent

Return type:

Training mode (dict)

classmethod from_pretrained_encoder_decoder(pretrained_path, strict=False, verbose=True, **kwargs)[source]

Create model and load encoder/decoder weights from pretrained checkpoint.

Factory method that creates a new model instance and loads encoder/decoder weights in one step. Useful for initializing new architectures from pretrained backbones.

Parameters:

  • pretrained_path (str | Path) – Path to pretrained weights

  • strict (bool, default: False) – If True, raise on shape mismatches

  • verbose (bool, default: True) – If True, print loading summary

  • **kwargs – Model configuration (passed to __init__)

Return type:

BaseTokenizer

Returns:

New model instance with loaded encoder/decoder weights

Example

>>> # Create FSQ model initialized from VAE weights
>>> model = DiscreteTokenizer.from_pretrained_encoder_decoder(
...     'weights/maisi_converted.pt',
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     z_channels=256,
...     channels=64,
...     channels_mult=(1, 2, 4),
... )
get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The buffer referenced by target

Return type:

torch.Tensor

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Any extra state to store in the module’s state_dict

Return type:

object

get_latent_shape(input_shape)[source]

Calculate output latent shape for given input shape.

Useful for pre-allocating memory or understanding compression ratio.

Parameters:

input_shape (tuple[int, ...]) – Input tensor shape (B, C, H, W) or (B, C, H, W, D)

Return type:

tuple[int, ...]

Returns:

Expected latent shape (B, latent_channels, H’, W’) or (B, C’, H’, W’, D’) where spatial dims are compressed by spatial_compression factor

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The Parameter referenced by target

Return type:

torch.nn.Parameter

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

Returns:

The submodule referenced by target

Return type:

torch.nn.Module

Raises:

AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

inference_mode()[source]

Context manager for optimized inference.

Configures the model for maximum inference performance: - Sets model to eval mode - Disables gradient computation - Uses torch.inference_mode for additional optimizations

The model is restored to its previous state upon exit.

Return type:

Generator[BaseTokenizer, None, None]

Example

>>> with model.inference_mode():
...     latents = model.tokenize(volume)
...     recon = model.detokenize(latents)
Yields:

self – The model instance for method chaining

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

load_encoder_decoder_weights(pretrained_path, strict=False, verbose=True)[source]

Load encoder/decoder weights from a pretrained checkpoint.

This method enables transfer learning by loading encoder and decoder weights from any compatible tokenizer (VAE, VQ-VAE, FSQ, etc.) while leaving quantizer-specific layers randomly initialized.

Use Cases

  1. Initialize discrete tokenizer from VAE: Load MAISI VAE encoder/decoder, then train only the quantizer layers (VQ, FSQ, etc.)

  2. Fine-tune on new domain: Start from pretrained weights, fine-tune all

  3. Encoder-only transfer: Use pretrained encoder for downstream tasks

How It Works

The method identifies encoder/decoder weights by key prefixes and loads only those that match. Quantizer-specific layers (quant_conv, post_quant_conv, quantizer) may or may not be loaded depending on architecture compatibility.

Weight Matching Strategy: - encoder.* keys: Always attempted - decoder.* keys: Always attempted - quant_conv.* keys: Loaded if shapes match - post_quant_conv.* keys: Loaded if shapes match - quantizer.* keys: Skipped (architecture-specific)

type pretrained_path:

str | Path

param pretrained_path:

Path to pretrained weights file (.pt, .bin) or directory containing β€˜pytorch_model.bin’

type strict:

bool, default: False

param strict:

If True, raise error on shape mismatches. If False (default), skip mismatched weights and log warnings.

type verbose:

bool, default: True

param verbose:

If True, print summary of loaded/skipped weights

rtype:

tuple[list[str], list[str]]

returns:

Tuple of (loaded_keys, skipped_keys) for inspection

Example

>>> # Initialize FSQ tokenizer from MAISI VAE weights
>>> model = DiscreteTokenizer(
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     # ... other args matching MAISI architecture
... )
>>> loaded, skipped = model.load_encoder_decoder_weights(
...     'weights/maisi_converted.pt',
...     verbose=True
... )
>>> print(f"Loaded {len(loaded)} keys, skipped {len(skipped)}")
>>>
>>> # Now train with frozen encoder if desired
>>> for param in model.encoder.parameters():
...     param.requires_grad = False

Note

For best results, ensure the pretrained model has the same: - channels, channels_mult, num_res_blocks - spatial_compression

The following can differ: - z_channels, latent_channels, embedding_dim - Quantizer type and configuration

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

  • state_dict (dict) – a dict containing parameters and persistent buffers.

  • strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

  • assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

  • missing_keys is a list of str containing any keys that are expected

    by this module but missing from the provided state_dict.

  • unexpected_keys is a list of str containing the keys that are not

    expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:

remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.

Yields:

Module – a module in the network

Return type:

Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

  • prefix (str) – prefix to prepend to all buffer names.

  • recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

  • remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Iterator[tuple[str, Tensor]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())
named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Return type:

Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

  • memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result

  • prefix (str, default: '') – a prefix that will be added to the name of the module

  • remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

  • prefix (str) – prefix to prepend to all parameter names.

  • recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

  • remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Iterator[tuple[str, Parameter]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())
num_parameters()[source]

Get total number of learnable parameters.

Return type:

int

Returns:

Sum of numel() for all parameters

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Return type:

Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
push_to_hub(repo_id, save_directory=None, commit_message='Upload model', private=False, **kwargs)[source]

Upload model to HuggingFace Hub.

Creates or updates a repository on the Hub with model weights and configuration.

Parameters:

  • repo_id (str) – Repository ID (e.g., β€œusername/model-name”)

  • save_directory (Optional[str], default: None) – Local directory to save before upload (auto-created if None)

  • commit_message (str, default: 'Upload model') – Commit message for the upload

  • private (bool, default: False) – Whether to create a private repository

  • **kwargs – Additional arguments for HfApi.upload_folder

Return type:

None

Example

>>> model.push_to_hub("username/my-vae-tokenizer")
>>> # Creates https://huggingface.co/username/my-vae-tokenizer
reconstruct(x, roi_size=None, overlap=0.5, sw_batch_size=1)[source]

Full encode-decode reconstruction with optional sliding window.

For large 3D volumes that exceed GPU memory, this method implements sliding window inference with Gaussian importance weighting to seamlessly blend overlapping patches.

The Algorithm

  1. Pad input to multiple of stride + window size

  2. Extract overlapping windows with specified stride

  3. Process windows in batches through tokenize -> detokenize

  4. Weight each window’s contribution by Gaussian importance

  5. Normalize by accumulated importance and crop to original size

Gaussian Weighting

Uses a Gaussian importance map (2D or 3D based on input) that gives higher weight to the center of each window, reducing boundary artifacts when blending.

type x:

Tensor

param x:

Input tensor of shape (B, C, H, W, D) for 3D or (B, C, H, W) for 2D

type roi_size:

Union[tuple[int, ...], int, None], default: None

param roi_size:

Size of sliding window. If None, processes entire volume. Can be int (isotropic) or tuple (anisotropic).

type overlap:

float, default: 0.5

param overlap:

Fraction of overlap between windows (0.0 to 0.9). Higher overlap = smoother blending but more compute.

type sw_batch_size:

int, default: 1

param sw_batch_size:

Number of windows to process in parallel per batch. Higher values use more GPU memory but are faster. Default is 1 (sequential processing).

rtype:

Tensor

returns:

Reconstructed tensor with same shape as input

Example

>>> volume = torch.randn(1, 1, 256, 256, 256)  # 256Β³ volume
>>> # Process in 128Β³ windows with 50% overlap, 4 windows at a time
>>> recon = model.reconstruct(volume, roi_size=128, overlap=0.5, sw_batch_size=4)

Note

  • For volumes that fit in memory, omit roi_size for faster processing

  • Overlap of 0.5 is a good default; higher values reduce artifacts but increase computation proportionally

  • sw_batch_size > 1 can significantly speed up inference on GPUs with sufficient memory

register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

  • name (str) – name of the buffer. The buffer can be accessed from this module using the given name

  • tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.

  • persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

None

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))
register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False

  • always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

  1. Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

  2. If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

  3. If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::

hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:

hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:
Return type:

None

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

  • name (str) – name of the parameter. The parameter can be accessed from this module using the given name

  • param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

None

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::

hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.

Returns:

self

Return type:

Module

save_pretrained(save_directory, push_to_hub=False, **kwargs)[source]

Save model weights and configuration to directory.

Creates a directory structure compatible with from_pretrained(): ` save_directory/ β”œβ”€β”€ config.json      # Model configuration └── pytorch_model.bin # Model weights `

Parameters:

  • save_directory (str | Path) – Path to save model

  • push_to_hub (bool, default: False) – If True, also upload to HuggingFace Hub

  • **kwargs – Additional arguments for push_to_hub

Return type:

None

Example

>>> model.save_pretrained("./my-tokenizer")
>>> # Later...
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict

Return type:

None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

  • target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

  • module (Module) – The module to set the submodule to.

  • strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

  • ValueError – If the target string is empty or if module is not an instance of nn.Module.

  • AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

None

share_memory()[source]

See torch.Tensor.share_memory_().

Return type:

Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

  • destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]
to(dtype, non_blocking=False)[source]
to(tensor, non_blocking=False)[source]
to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

  • device (torch.device) – the desired device of the parameters and buffers in this module

  • dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

  • tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

  • memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

  • device (torch.device) – The desired device of the parameters and buffers in this module.

  • recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

tokenize(x)[source]

Encode input to latent representation (convenience method).

For inference, this is the primary encoding method. For VAE, returns the sampled latent (not the mean), enabling diverse reconstructions.

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input tensor

Return type:

Float[Tensor, 'batch latent_channels *spatial_compressed']

Returns:

Latent tensor suitable for storage, manipulation, or decoding

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:

dst_type (type or string) – the desired type

Returns:

self

Return type:

Module

weights_name: str = 'pytorch_model.bin'
xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:

set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.

Return type:

None

config: dict[str, Any]
training: bool
class medtokenizers.TiTokTokenizer(dim, in_channels=1, out_channels=None, num_tokens=32, num_embeddings=1024, embedding_dim=None, hidden_dim=256, num_heads=8, num_layers=4, patch_size=16, resolution=128, dropout=0.0, beta=0.25, use_norm=False, use_ema=False, ema_decay=0.99, reset_unused_codes=False, dead_code_threshold=100, name='TiTokTokenizer')[source][source]

Bases: medtokenizers.modules.base.BaseTokenizer

Transformer-based 1D image tokenizer (TiTok).

Implements the TiTok idea (Yu et al., 2024) for 2D and 3D medical images: instead of producing a spatial grid of codes, the image is compressed into a flat sequence of num_tokens learnable 1D latent tokens. The input is split into non-overlapping patches and embedded; a learnable set of latent tokens is appended and a Transformer encoder lets those latent tokens attend over all patches. The latent tokens are then vector-quantized into discrete indices. Decoding mirrors this: quantized latents plus a learnable mask token per patch are passed through a Transformer decoder, and the patch outputs are un-embedded and reassembled into an image.

Because the latent is a fixed-length 1D sequence (independent of spatial resolution), TiTok is well suited to autoregressive / sequence models. Note that the input spatial size is fixed: it must equal resolution (validated at encode time), and resolution must be divisible by patch_size.

Patch ordering:

  • 2D: patches are flattened in (height, width) row-major order, i.e. token index i = row * grid_w + col. Each patch vector concatenates its pixels as (p_h, p_w, channels).

  • 3D: patches are flattened in (depth, height, width) order following the (g_h, g_w, g_d) grid, with each patch vector laid out as (p_h, p_w, p_d, channels).

_patchify / _unpatchify are exact inverses, so decode restores the original spatial layout.

Parameters:

  • dim (int) – Spatial dimensionality, 2 or 3.

  • in_channels (int, default: 1) – Number of input image channels.

  • out_channels (int | None, default: None) – Number of reconstructed output channels. Defaults to in_channels when None.

  • num_tokens (int, default: 32) – Number of 1D latent tokens (the compressed sequence length).

  • num_embeddings (int, default: 1024) – Codebook size of the vector quantizer (vocabulary).

  • embedding_dim (int | None, default: None) – Dimensionality of each quantized code. Defaults to hidden_dim when None.

  • hidden_dim (int, default: 256) – Transformer model width; must be divisible by num_heads.

  • num_heads (int, default: 8) – Number of attention heads in encoder and decoder.

  • num_layers (int, default: 4) – Number of Transformer layers in encoder and decoder.

  • patch_size (Union[int, Iterable[int]], default: 16) – Patch edge length. Either a single int (applied to every spatial axis) or a per-axis iterable of length dim.

  • resolution (Union[int, Iterable[int]], default: 128) – Expected input spatial size. Either a single int or a per-axis iterable of length dim. Must be divisible by patch_size along every axis.

  • dropout (float, default: 0.0) – Dropout probability in [0, 1) applied to embeddings and within Transformer layers.

  • beta (float, default: 0.25) – Commitment loss weight for the vector quantizer.

  • use_norm (bool, default: False) – Whether the quantizer L2-normalizes codes/inputs.

  • use_ema (bool, default: False) – Whether the quantizer updates its codebook via EMA.

  • ema_decay (float, default: 0.99) – EMA decay used when use_ema is True.

  • reset_unused_codes (bool, default: False) – Whether to reset dead codebook entries.

  • dead_code_threshold (int, default: 100) – Usage count below which a code is considered dead.

  • name (str, default: 'TiTokTokenizer') – Human-readable tokenizer name.

Shapes:

  • Input: (B, in_channels, *resolution) where *resolution is (H, W) for 2D or (H, W, D) for 3D.

  • Discrete indices: (B, num_tokens) (integer dtype).

  • Quantized latents: (B, num_tokens, embedding_dim).

  • Reconstruction: (B, out_channels, *resolution).

__init__(dim, in_channels=1, out_channels=None, num_tokens=32, num_embeddings=1024, embedding_dim=None, hidden_dim=256, num_heads=8, num_layers=4, patch_size=16, resolution=128, dropout=0.0, beta=0.25, use_norm=False, use_ema=False, ema_decay=0.99, reset_unused_codes=False, dead_code_threshold=100, name='TiTokTokenizer')[source][source]
Parameters:

  • dim (int)

  • in_channels (int, default: 1)

  • out_channels (int | None, default: None)

  • num_tokens (int, default: 32)

  • num_embeddings (int, default: 1024)

  • embedding_dim (int | None, default: None)

  • hidden_dim (int, default: 256)

  • num_heads (int, default: 8)

  • num_layers (int, default: 4)

  • patch_size (Union[int, Iterable[int]], default: 16)

  • resolution (Union[int, Iterable[int]], default: 128)

  • dropout (float, default: 0.0)

  • beta (float, default: 0.25)

  • use_norm (bool, default: False)

  • use_ema (bool, default: False)

  • ema_decay (float, default: 0.99)

  • reset_unused_codes (bool, default: False)

  • dead_code_threshold (int, default: 100)

  • name (str, default: 'TiTokTokenizer')

config: dict[str, Any]
encode(x)[source][source]

Encode an image into discrete 1D latent tokens.

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input image of shape (B, in_channels, *resolution).

Return type:

tuple[Int[Tensor, 'batch num_tokens'], Float[Tensor, 'batch num_tokens embedding_dim'], Tensor]

Returns:

Tuple of (indices, quantized, quant_loss) where indices has shape (B, num_tokens) (integer dtype), quantized has shape (B, num_tokens, embedding_dim), and quant_loss is the quantizer commitment/codebook loss.

indices_to_codes(indices)[source][source]

Look up quantized code vectors for a batch of token indices.

Parameters:

indices (Int[Tensor, 'batch num_tokens']) – Integer indices of shape (B, num_tokens).

Return type:

Float[Tensor, 'batch num_tokens embedding_dim']

Returns:

Quantized code vectors of shape (B, num_tokens, embedding_dim).

decode(quantized)[source][source]

Decode quantized 1D latent tokens back into an image.

Parameters:

quantized (Float[Tensor, 'batch num_tokens embedding_dim']) – Quantized code vectors of shape (B, num_tokens, embedding_dim).

Return type:

Float[Tensor, 'batch channels *spatial']

Returns:

Reconstructed image of shape (B, out_channels, *resolution).

forward(x)[source][source]

Encode then decode x, returning training or eval outputs.

Parameters:

x (Float[Tensor, 'batch channels *spatial']) – Input image of shape (B, in_channels, *resolution).

Return type:

dict[str, Tensor] | NetworkEval

Returns:

In training mode, a dict with reconstructions, quant_loss, quant_info (indices) and latents (quantized). In eval mode, a NetworkEval named tuple.

tokenize(x)[source][source]

Encode x to discrete token indices of shape (B, num_tokens).

Parameters:

x (Float[Tensor, 'batch channels *spatial'])

Return type:

Int[Tensor, 'batch num_tokens']

detokenize(indices)[source][source]

Decode discrete token indices back into an image.

Inverse of tokenize(). This is the canonical index-to-reconstruction decoding path.

Parameters:

indices (Int[Tensor, 'batch num_tokens']) – Integer indices of shape (B, num_tokens).

Return type:

Float[Tensor, 'batch channels *spatial']

Returns:

Reconstructed image of shape (B, out_channels, *resolution).

T_destination = ~T_destination
add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

  • name (str) – name of the child module. The child module can be accessed from this module using the given name

  • module (Module) – child module to be added to the module.

Return type:

None

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:

fn (Module -> None) – function to be applied to each submodule

Returns:

self

Return type:

Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor – module buffer

Return type:

Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
call_super_init: bool = False
children()[source]

Return an iterator over immediate children modules.

Yields:

Module – a child module

Return type:

Iterator[Module]

compile(mode='reduce-overhead', fullgraph=False, **kwargs)[source]

Compile model with torch.compile() for faster inference.

Uses PyTorch 2.0+ compilation to optimize the model graph. The compiled model maintains the same interface but runs faster, especially for repeated inference calls.

Parameters:

  • mode (str, default: 'reduce-overhead') – Compilation mode. Options: - β€œreduce-overhead”: Best for small batches (default) - β€œmax-autotune”: Best throughput, longer warmup - β€œdefault”: Balanced compilation

  • fullgraph (bool, default: False) – If True, require full graph compilation (stricter)

  • **kwargs – Additional arguments passed to torch.compile()

Return type:

BaseTokenizer

Returns:

Compiled model (self, for method chaining)

Example

>>> model = ContinuousTokenizer.from_pretrained("path/to/model")
>>> model = model.compile(mode="reduce-overhead")
>>> # First call triggers compilation (slower)
>>> latents = model.tokenize(batch)
>>> # Subsequent calls are faster
>>> latents = model.tokenize(batch2)

Note

  • Compilation happens on first forward pass

  • Different input shapes may trigger recompilation

  • Use dynamic=False (default) for fixed input sizes

config_name: str = 'config.json'
cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

decode_batch(z, batch_size=8, show_progress=False)[source]

Decode large batch with automatic mini-batching.

Parameters:

  • z (Tensor) – Full latent tensor

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated reconstructions for all inputs

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

dump_patches: bool = False
encode_batch(x, batch_size=8, show_progress=False)[source]

Encode large batch with automatic mini-batching.

Processes input in chunks to avoid OOM for large datasets.

Parameters:

  • x (Tensor) – Full input tensor of shape (N, C, *spatial)

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated latents for all inputs

Example

>>> dataset = torch.randn(1000, 1, 64, 64, 64)
>>> latents = model.encode_batch(dataset, batch_size=4)
eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

self

Return type:

Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

classmethod from_pretrained(model_name_or_path, map_location=None, **kwargs)[source]

Load model from local directory or HuggingFace Hub.

Automatically detects whether path is local or a Hub repository. For Hub repos, downloads config and weights to cache.

Parameters:

  • model_name_or_path (str) – Local path or HuggingFace Hub repo ID

  • map_location (Optional[str], default: None) – Device to load weights to (default: auto-detect)

  • **kwargs – Override config parameters

Return type:

BaseTokenizer

Returns:

Loaded model instance

Example

>>> # From local path
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
>>>
>>> # From HuggingFace Hub
>>> model = ContinuousTokenizer.from_pretrained("username/my-tokenizer")
>>>
>>> # Override config
>>> model = ContinuousTokenizer.from_pretrained(
...     "./my-tokenizer",
...     dropout=0.1  # Override saved dropout value
... )
classmethod from_pretrained_encoder_decoder(pretrained_path, strict=False, verbose=True, **kwargs)[source]

Create model and load encoder/decoder weights from pretrained checkpoint.

Factory method that creates a new model instance and loads encoder/decoder weights in one step. Useful for initializing new architectures from pretrained backbones.

Parameters:

  • pretrained_path (str | Path) – Path to pretrained weights

  • strict (bool, default: False) – If True, raise on shape mismatches

  • verbose (bool, default: True) – If True, print loading summary

  • **kwargs – Model configuration (passed to __init__)

Return type:

BaseTokenizer

Returns:

New model instance with loaded encoder/decoder weights

Example

>>> # Create FSQ model initialized from VAE weights
>>> model = DiscreteTokenizer.from_pretrained_encoder_decoder(
...     'weights/maisi_converted.pt',
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     z_channels=256,
...     channels=64,
...     channels_mult=(1, 2, 4),
... )
get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The buffer referenced by target

Return type:

torch.Tensor

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Any extra state to store in the module’s state_dict

Return type:

object

get_latent_shape(input_shape)[source]

Calculate latent shape for given input shape.

Parameters:

input_shape (tuple[int, ...]) – Input tensor shape (B, C, *spatial)

Return type:

tuple[int, ...]

Returns:

Expected latent tensor shape

Raises:

NotImplementedError – Must be implemented by subclass

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The Parameter referenced by target

Return type:

torch.nn.Parameter

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

Returns:

The submodule referenced by target

Return type:

torch.nn.Module

Raises:

AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

inference_mode()[source]

Context manager for optimized inference.

Configures the model for maximum inference performance: - Sets model to eval mode - Disables gradient computation - Uses torch.inference_mode for additional optimizations

The model is restored to its previous state upon exit.

Return type:

Generator[BaseTokenizer, None, None]

Example

>>> with model.inference_mode():
...     latents = model.tokenize(volume)
...     recon = model.detokenize(latents)
Yields:

self – The model instance for method chaining

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

load_encoder_decoder_weights(pretrained_path, strict=False, verbose=True)[source]

Load encoder/decoder weights from a pretrained checkpoint.

This method enables transfer learning by loading encoder and decoder weights from any compatible tokenizer (VAE, VQ-VAE, FSQ, etc.) while leaving quantizer-specific layers randomly initialized.

Use Cases

  1. Initialize discrete tokenizer from VAE: Load MAISI VAE encoder/decoder, then train only the quantizer layers (VQ, FSQ, etc.)

  2. Fine-tune on new domain: Start from pretrained weights, fine-tune all

  3. Encoder-only transfer: Use pretrained encoder for downstream tasks

How It Works

The method identifies encoder/decoder weights by key prefixes and loads only those that match. Quantizer-specific layers (quant_conv, post_quant_conv, quantizer) may or may not be loaded depending on architecture compatibility.

Weight Matching Strategy: - encoder.* keys: Always attempted - decoder.* keys: Always attempted - quant_conv.* keys: Loaded if shapes match - post_quant_conv.* keys: Loaded if shapes match - quantizer.* keys: Skipped (architecture-specific)

type pretrained_path:

str | Path

param pretrained_path:

Path to pretrained weights file (.pt, .bin) or directory containing β€˜pytorch_model.bin’

type strict:

bool, default: False

param strict:

If True, raise error on shape mismatches. If False (default), skip mismatched weights and log warnings.

type verbose:

bool, default: True

param verbose:

If True, print summary of loaded/skipped weights

rtype:

tuple[list[str], list[str]]

returns:

Tuple of (loaded_keys, skipped_keys) for inspection

Example

>>> # Initialize FSQ tokenizer from MAISI VAE weights
>>> model = DiscreteTokenizer(
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     # ... other args matching MAISI architecture
... )
>>> loaded, skipped = model.load_encoder_decoder_weights(
...     'weights/maisi_converted.pt',
...     verbose=True
... )
>>> print(f"Loaded {len(loaded)} keys, skipped {len(skipped)}")
>>>
>>> # Now train with frozen encoder if desired
>>> for param in model.encoder.parameters():
...     param.requires_grad = False

Note

For best results, ensure the pretrained model has the same: - channels, channels_mult, num_res_blocks - spatial_compression

The following can differ: - z_channels, latent_channels, embedding_dim - Quantizer type and configuration

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

  • state_dict (dict) – a dict containing parameters and persistent buffers.

  • strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

  • assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

  • missing_keys is a list of str containing any keys that are expected

    by this module but missing from the provided state_dict.

  • unexpected_keys is a list of str containing the keys that are not

    expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:

remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.

Yields:

Module – a module in the network

Return type:

Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

  • prefix (str) – prefix to prepend to all buffer names.

  • recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

  • remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Iterator[tuple[str, Tensor]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())
named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Return type:

Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

  • memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result

  • prefix (str, default: '') – a prefix that will be added to the name of the module

  • remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

  • prefix (str) – prefix to prepend to all parameter names.

  • recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

  • remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Iterator[tuple[str, Parameter]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())
num_parameters()[source]

Get total number of learnable parameters.

Return type:

int

Returns:

Sum of numel() for all parameters

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Return type:

Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
push_to_hub(repo_id, save_directory=None, commit_message='Upload model', private=False, **kwargs)[source]

Upload model to HuggingFace Hub.

Creates or updates a repository on the Hub with model weights and configuration.

Parameters:

  • repo_id (str) – Repository ID (e.g., β€œusername/model-name”)

  • save_directory (Optional[str], default: None) – Local directory to save before upload (auto-created if None)

  • commit_message (str, default: 'Upload model') – Commit message for the upload

  • private (bool, default: False) – Whether to create a private repository

  • **kwargs – Additional arguments for HfApi.upload_folder

Return type:

None

Example

>>> model.push_to_hub("username/my-vae-tokenizer")
>>> # Creates https://huggingface.co/username/my-vae-tokenizer
reconstruct(x, roi_size=None, overlap=0.5, sw_batch_size=1)[source]

Full encode-decode reconstruction with optional sliding window.

For large 3D volumes that exceed GPU memory, this method implements sliding window inference with Gaussian importance weighting to seamlessly blend overlapping patches.

The Algorithm

  1. Pad input to multiple of stride + window size

  2. Extract overlapping windows with specified stride

  3. Process windows in batches through tokenize -> detokenize

  4. Weight each window’s contribution by Gaussian importance

  5. Normalize by accumulated importance and crop to original size

Gaussian Weighting

Uses a Gaussian importance map (2D or 3D based on input) that gives higher weight to the center of each window, reducing boundary artifacts when blending.

type x:

Tensor

param x:

Input tensor of shape (B, C, H, W, D) for 3D or (B, C, H, W) for 2D

type roi_size:

Union[tuple[int, ...], int, None], default: None

param roi_size:

Size of sliding window. If None, processes entire volume. Can be int (isotropic) or tuple (anisotropic).

type overlap:

float, default: 0.5

param overlap:

Fraction of overlap between windows (0.0 to 0.9). Higher overlap = smoother blending but more compute.

type sw_batch_size:

int, default: 1

param sw_batch_size:

Number of windows to process in parallel per batch. Higher values use more GPU memory but are faster. Default is 1 (sequential processing).

rtype:

Tensor

returns:

Reconstructed tensor with same shape as input

Example

>>> volume = torch.randn(1, 1, 256, 256, 256)  # 256Β³ volume
>>> # Process in 128Β³ windows with 50% overlap, 4 windows at a time
>>> recon = model.reconstruct(volume, roi_size=128, overlap=0.5, sw_batch_size=4)

Note

  • For volumes that fit in memory, omit roi_size for faster processing

  • Overlap of 0.5 is a good default; higher values reduce artifacts but increase computation proportionally

  • sw_batch_size > 1 can significantly speed up inference on GPUs with sufficient memory

register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

  • name (str) – name of the buffer. The buffer can be accessed from this module using the given name

  • tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.

  • persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

None

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))
register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False

  • always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

  1. Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

  2. If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

  3. If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::

hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:

hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:
Return type:

None

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

  • name (str) – name of the parameter. The parameter can be accessed from this module using the given name

  • param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

None

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::

hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.

Returns:

self

Return type:

Module

save_pretrained(save_directory, push_to_hub=False, **kwargs)[source]

Save model weights and configuration to directory.

Creates a directory structure compatible with from_pretrained(): ` save_directory/ β”œβ”€β”€ config.json      # Model configuration └── pytorch_model.bin # Model weights `

Parameters:

  • save_directory (str | Path) – Path to save model

  • push_to_hub (bool, default: False) – If True, also upload to HuggingFace Hub

  • **kwargs – Additional arguments for push_to_hub

Return type:

None

Example

>>> model.save_pretrained("./my-tokenizer")
>>> # Later...
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict

Return type:

None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

  • target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

  • module (Module) – The module to set the submodule to.

  • strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

  • ValueError – If the target string is empty or if module is not an instance of nn.Module.

  • AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

None

share_memory()[source]

See torch.Tensor.share_memory_().

Return type:

Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

  • destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]
to(dtype, non_blocking=False)[source]
to(tensor, non_blocking=False)[source]
to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

  • device (torch.device) – the desired device of the parameters and buffers in this module

  • dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

  • tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

  • memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

  • device (torch.device) – The desired device of the parameters and buffers in this module.

  • recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:

dst_type (type or string) – the desired type

Returns:

self

Return type:

Module

weights_name: str = 'pytorch_model.bin'
xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:

set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.

Return type:

None

training: bool
class medtokenizers.RAETokenizer(dim, encoder_type, encoder_name_or_path, out_channels=1, latent_dim=None, patch_size=None, encoder_image_size=None, encoder_drop_cls_token=True, encoder_kwargs=None, decoder_hidden_dim=1024, decoder_num_layers=2, decoder_dropout=0.0, noise_tau=0.0, latent_stats_path=None, latent_eps=1e-05, name='RAETokenizer')[source][source]

Bases: medtokenizers.modules.base.BaseTokenizer

Representation Autoencoder (RAE) tokenizer with a frozen encoder.

Parameters:

  • dim (int)

  • encoder_type (Literal['vit', 'medsiglip', 'neurovfm'])

  • encoder_name_or_path (str)

  • out_channels (int, default: 1)

  • latent_dim (int | None, default: None)

  • patch_size (int | tuple[int, ...] | None, default: None)

  • encoder_image_size (int | None, default: None)

  • encoder_drop_cls_token (bool, default: True)

  • encoder_kwargs (dict[str, Any] | None, default: None)

  • decoder_hidden_dim (int, default: 1024)

  • decoder_num_layers (int, default: 2)

  • decoder_dropout (float, default: 0.0)

  • noise_tau (float, default: 0.0)

  • latent_stats_path (str | None, default: None)

  • latent_eps (float, default: 1e-05)

  • name (str, default: 'RAETokenizer')

__init__(dim, encoder_type, encoder_name_or_path, out_channels=1, latent_dim=None, patch_size=None, encoder_image_size=None, encoder_drop_cls_token=True, encoder_kwargs=None, decoder_hidden_dim=1024, decoder_num_layers=2, decoder_dropout=0.0, noise_tau=0.0, latent_stats_path=None, latent_eps=1e-05, name='RAETokenizer')[source][source]
Parameters:

  • dim (int)

  • encoder_type (Literal['vit', 'medsiglip', 'neurovfm'])

  • encoder_name_or_path (str)

  • out_channels (int, default: 1)

  • latent_dim (int | None, default: None)

  • patch_size (int | tuple[int, ...] | None, default: None)

  • encoder_image_size (int | None, default: None)

  • encoder_drop_cls_token (bool, default: True)

  • encoder_kwargs (dict[str, Any] | None, default: None)

  • decoder_hidden_dim (int, default: 1024)

  • decoder_num_layers (int, default: 2)

  • decoder_dropout (float, default: 0.0)

  • noise_tau (float, default: 0.0)

  • latent_stats_path (str | None, default: None)

  • latent_eps (float, default: 1e-05)

  • name (str, default: 'RAETokenizer')

config: dict[str, Any]
T_destination = ~T_destination
add_module(name, module)[source]

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

  • name (str) – name of the child module. The child module can be accessed from this module using the given name

  • module (Module) – child module to be added to the module.

Return type:

None

apply(fn)[source]

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Parameters:

fn (Module -> None) – function to be applied to each submodule

Returns:

self

Return type:

Module

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
bfloat16()[source]

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

buffers(recurse=True)[source]

Return an iterator over module buffers.

Parameters:

recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor – module buffer

Return type:

Iterator[Tensor]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
call_super_init: bool = False
children()[source]

Return an iterator over immediate children modules.

Yields:

Module – a child module

Return type:

Iterator[Module]

compile(mode='reduce-overhead', fullgraph=False, **kwargs)[source]

Compile model with torch.compile() for faster inference.

Uses PyTorch 2.0+ compilation to optimize the model graph. The compiled model maintains the same interface but runs faster, especially for repeated inference calls.

Parameters:

  • mode (str, default: 'reduce-overhead') – Compilation mode. Options: - β€œreduce-overhead”: Best for small batches (default) - β€œmax-autotune”: Best throughput, longer warmup - β€œdefault”: Balanced compilation

  • fullgraph (bool, default: False) – If True, require full graph compilation (stricter)

  • **kwargs – Additional arguments passed to torch.compile()

Return type:

BaseTokenizer

Returns:

Compiled model (self, for method chaining)

Example

>>> model = ContinuousTokenizer.from_pretrained("path/to/model")
>>> model = model.compile(mode="reduce-overhead")
>>> # First call triggers compilation (slower)
>>> latents = model.tokenize(batch)
>>> # Subsequent calls are faster
>>> latents = model.tokenize(batch2)

Note

  • Compilation happens on first forward pass

  • Different input shapes may trigger recompilation

  • Use dynamic=False (default) for fixed input sizes

config_name: str = 'config.json'
cpu()[source]

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

cuda(device=None)[source]

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

decode_batch(z, batch_size=8, show_progress=False)[source]

Decode large batch with automatic mini-batching.

Parameters:

  • z (Tensor) – Full latent tensor

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated reconstructions for all inputs

detokenize(z)[source]

Decode from latent space (convenience method).

Parameters:

z (Tensor) – Latent tensor

Return type:

Tensor

Returns:

Reconstructed output

double()[source]

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

dump_patches: bool = False
encode_batch(x, batch_size=8, show_progress=False)[source]

Encode large batch with automatic mini-batching.

Processes input in chunks to avoid OOM for large datasets.

Parameters:

  • x (Tensor) – Full input tensor of shape (N, C, *spatial)

  • batch_size (int, default: 8) – Mini-batch size for processing

  • show_progress (bool, default: False) – Whether to show tqdm progress bar

Return type:

Tensor

Returns:

Concatenated latents for all inputs

Example

>>> dataset = torch.randn(1000, 1, 64, 64, 64)
>>> latents = model.encode_batch(dataset, batch_size=4)
eval()[source]

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

self

Return type:

Module

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

float()[source]

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

classmethod from_pretrained(model_name_or_path, map_location=None, **kwargs)[source]

Load model from local directory or HuggingFace Hub.

Automatically detects whether path is local or a Hub repository. For Hub repos, downloads config and weights to cache.

Parameters:

  • model_name_or_path (str) – Local path or HuggingFace Hub repo ID

  • map_location (Optional[str], default: None) – Device to load weights to (default: auto-detect)

  • **kwargs – Override config parameters

Return type:

BaseTokenizer

Returns:

Loaded model instance

Example

>>> # From local path
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
>>>
>>> # From HuggingFace Hub
>>> model = ContinuousTokenizer.from_pretrained("username/my-tokenizer")
>>>
>>> # Override config
>>> model = ContinuousTokenizer.from_pretrained(
...     "./my-tokenizer",
...     dropout=0.1  # Override saved dropout value
... )
classmethod from_pretrained_encoder_decoder(pretrained_path, strict=False, verbose=True, **kwargs)[source]

Create model and load encoder/decoder weights from pretrained checkpoint.

Factory method that creates a new model instance and loads encoder/decoder weights in one step. Useful for initializing new architectures from pretrained backbones.

Parameters:

  • pretrained_path (str | Path) – Path to pretrained weights

  • strict (bool, default: False) – If True, raise on shape mismatches

  • verbose (bool, default: True) – If True, print loading summary

  • **kwargs – Model configuration (passed to __init__)

Return type:

BaseTokenizer

Returns:

New model instance with loaded encoder/decoder weights

Example

>>> # Create FSQ model initialized from VAE weights
>>> model = DiscreteTokenizer.from_pretrained_encoder_decoder(
...     'weights/maisi_converted.pt',
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     z_channels=256,
...     channels=64,
...     channels_mult=(1, 2, 4),
... )
get_buffer(target)[source]

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The buffer referenced by target

Return type:

torch.Tensor

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Any extra state to store in the module’s state_dict

Return type:

object

get_latent_shape(input_shape)[source]

Calculate latent shape for given input shape.

Parameters:

input_shape (tuple[int, ...]) – Input tensor shape (B, C, *spatial)

Return type:

tuple[int, ...]

Returns:

Expected latent tensor shape

Raises:

NotImplementedError – Must be implemented by subclass

get_parameter(target)[source]

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Parameters:

target (str) – The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

The Parameter referenced by target

Return type:

torch.nn.Parameter

Raises:

AttributeError – If the target string references an invalid path or resolves to something that is not an nn.Parameter

get_submodule(target)[source]

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

Returns:

The submodule referenced by target

Return type:

torch.nn.Module

Raises:

AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half()[source]

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:

self

Return type:

Module

inference_mode()[source]

Context manager for optimized inference.

Configures the model for maximum inference performance: - Sets model to eval mode - Disables gradient computation - Uses torch.inference_mode for additional optimizations

The model is restored to its previous state upon exit.

Return type:

Generator[BaseTokenizer, None, None]

Example

>>> with model.inference_mode():
...     latents = model.tokenize(volume)
...     recon = model.detokenize(latents)
Yields:

self – The model instance for method chaining

ipu(device=None)[source]

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

load_encoder_decoder_weights(pretrained_path, strict=False, verbose=True)[source]

Load encoder/decoder weights from a pretrained checkpoint.

This method enables transfer learning by loading encoder and decoder weights from any compatible tokenizer (VAE, VQ-VAE, FSQ, etc.) while leaving quantizer-specific layers randomly initialized.

Use Cases

  1. Initialize discrete tokenizer from VAE: Load MAISI VAE encoder/decoder, then train only the quantizer layers (VQ, FSQ, etc.)

  2. Fine-tune on new domain: Start from pretrained weights, fine-tune all

  3. Encoder-only transfer: Use pretrained encoder for downstream tasks

How It Works

The method identifies encoder/decoder weights by key prefixes and loads only those that match. Quantizer-specific layers (quant_conv, post_quant_conv, quantizer) may or may not be loaded depending on architecture compatibility.

Weight Matching Strategy: - encoder.* keys: Always attempted - decoder.* keys: Always attempted - quant_conv.* keys: Loaded if shapes match - post_quant_conv.* keys: Loaded if shapes match - quantizer.* keys: Skipped (architecture-specific)

type pretrained_path:

str | Path

param pretrained_path:

Path to pretrained weights file (.pt, .bin) or directory containing β€˜pytorch_model.bin’

type strict:

bool, default: False

param strict:

If True, raise error on shape mismatches. If False (default), skip mismatched weights and log warnings.

type verbose:

bool, default: True

param verbose:

If True, print summary of loaded/skipped weights

rtype:

tuple[list[str], list[str]]

returns:

Tuple of (loaded_keys, skipped_keys) for inspection

Example

>>> # Initialize FSQ tokenizer from MAISI VAE weights
>>> model = DiscreteTokenizer(
...     dim=3,
...     quantizer='FSQ',
...     levels=[8, 5, 5, 5],
...     # ... other args matching MAISI architecture
... )
>>> loaded, skipped = model.load_encoder_decoder_weights(
...     'weights/maisi_converted.pt',
...     verbose=True
... )
>>> print(f"Loaded {len(loaded)} keys, skipped {len(skipped)}")
>>>
>>> # Now train with frozen encoder if desired
>>> for param in model.encoder.parameters():
...     param.requires_grad = False

Note

For best results, ensure the pretrained model has the same: - channels, channels_mult, num_res_blocks - spatial_compression

The following can differ: - z_channels, latent_channels, embedding_dim - Quantizer type and configuration

load_state_dict(state_dict, strict=True, assign=False)[source]

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

  • state_dict (dict) – a dict containing parameters and persistent buffers.

  • strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

  • assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

  • missing_keys is a list of str containing any keys that are expected

    by this module but missing from the provided state_dict.

  • unexpected_keys is a list of str containing the keys that are not

    expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

modules(remove_duplicate=True)[source]

Return an iterator over all modules in the network.

Parameters:

remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not.

Yields:

Module – a module in the network

Return type:

Iterator[Module]

Note

Duplicate modules are returned only once by default. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
mtia(device=None)[source]

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

named_buffers(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

  • prefix (str) – prefix to prepend to all buffer names.

  • recurse (bool, optional) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

  • remove_duplicate (bool, optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor) – Tuple containing the name and buffer

Return type:

Iterator[tuple[str, Tensor]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())
named_children()[source]

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

(str, Module) – Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Return type:

Iterator[tuple[str, Module]]

named_modules(memo=None, prefix='', remove_duplicate=True)[source]

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

  • memo (set[Module] | None, default: None) – a memo to store the set of modules already added to the result

  • prefix (str, default: '') – a prefix that will be added to the name of the module

  • remove_duplicate (bool, default: True) – whether to remove the duplicated module instances in the result or not

Yields:

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
named_parameters(prefix='', recurse=True, remove_duplicate=True)[source]

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

  • prefix (str) – prefix to prepend to all parameter names.

  • recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

  • remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Return type:

Iterator[tuple[str, Parameter]]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())
num_parameters()[source]

Get total number of learnable parameters.

Return type:

int

Returns:

Sum of numel() for all parameters

parameters(recurse=True)[source]

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Return type:

Iterator[Parameter]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
push_to_hub(repo_id, save_directory=None, commit_message='Upload model', private=False, **kwargs)[source]

Upload model to HuggingFace Hub.

Creates or updates a repository on the Hub with model weights and configuration.

Parameters:

  • repo_id (str) – Repository ID (e.g., β€œusername/model-name”)

  • save_directory (Optional[str], default: None) – Local directory to save before upload (auto-created if None)

  • commit_message (str, default: 'Upload model') – Commit message for the upload

  • private (bool, default: False) – Whether to create a private repository

  • **kwargs – Additional arguments for HfApi.upload_folder

Return type:

None

Example

>>> model.push_to_hub("username/my-vae-tokenizer")
>>> # Creates https://huggingface.co/username/my-vae-tokenizer
reconstruct(x, roi_size=None, overlap=0.5, sw_batch_size=1)[source]

Full encode-decode reconstruction with optional sliding window.

For large 3D volumes that exceed GPU memory, this method implements sliding window inference with Gaussian importance weighting to seamlessly blend overlapping patches.

The Algorithm

  1. Pad input to multiple of stride + window size

  2. Extract overlapping windows with specified stride

  3. Process windows in batches through tokenize -> detokenize

  4. Weight each window’s contribution by Gaussian importance

  5. Normalize by accumulated importance and crop to original size

Gaussian Weighting

Uses a Gaussian importance map (2D or 3D based on input) that gives higher weight to the center of each window, reducing boundary artifacts when blending.

type x:

Tensor

param x:

Input tensor of shape (B, C, H, W, D) for 3D or (B, C, H, W) for 2D

type roi_size:

Union[tuple[int, ...], int, None], default: None

param roi_size:

Size of sliding window. If None, processes entire volume. Can be int (isotropic) or tuple (anisotropic).

type overlap:

float, default: 0.5

param overlap:

Fraction of overlap between windows (0.0 to 0.9). Higher overlap = smoother blending but more compute.

type sw_batch_size:

int, default: 1

param sw_batch_size:

Number of windows to process in parallel per batch. Higher values use more GPU memory but are faster. Default is 1 (sequential processing).

rtype:

Tensor

returns:

Reconstructed tensor with same shape as input

Example

>>> volume = torch.randn(1, 1, 256, 256, 256)  # 256Β³ volume
>>> # Process in 128Β³ windows with 50% overlap, 4 windows at a time
>>> recon = model.reconstruct(volume, roi_size=128, overlap=0.5, sw_batch_size=4)

Note

  • For volumes that fit in memory, omit roi_size for faster processing

  • Overlap of 0.5 is a good default; higher values reduce artifacts but increase computation proportionally

  • sw_batch_size > 1 can significantly speed up inference on GPUs with sufficient memory

register_backward_hook(hook)[source]

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], tuple[Tensor, ...] | Tensor | None])

register_buffer(name, tensor, persistent=True)[source]

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

  • name (str) – name of the buffer. The buffer can be accessed from this module using the given name

  • tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.

  • persistent (bool) – whether the buffer is part of this module’s state_dict.

Return type:

None

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))
register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)[source]

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False

  • always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)[source]

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs
Parameters:

  • hook (Callable) – The user defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

  • with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook, prepend=False)[source]

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

  1. Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

  2. If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

  3. If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)[source]

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Parameters:

  • hook (Callable) – The user-defined hook to be registered.

  • prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::

hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters:

hook (Callable) – Callable hook that will be invoked before loading the state dict.

register_module(name, module)[source]

Alias for add_module().

Parameters:
Return type:

None

register_parameter(name, param)[source]

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

  • name (str) – name of the parameter. The parameter can be accessed from this module using the given name

  • param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Return type:

None

register_state_dict_post_hook(hook)[source]

Register a post-hook for the state_dict() method.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)[source]

Register a pre-hook for the state_dict() method.

It should have the following signature::

hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)[source]

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

requires_grad (bool) – whether autograd should record operations on parameters in this module. Default: True.

Returns:

self

Return type:

Module

save_pretrained(save_directory, push_to_hub=False, **kwargs)[source]

Save model weights and configuration to directory.

Creates a directory structure compatible with from_pretrained(): ` save_directory/ β”œβ”€β”€ config.json      # Model configuration └── pytorch_model.bin # Model weights `

Parameters:

  • save_directory (str | Path) – Path to save model

  • push_to_hub (bool, default: False) – If True, also upload to HuggingFace Hub

  • **kwargs – Additional arguments for push_to_hub

Return type:

None

Example

>>> model.save_pretrained("./my-tokenizer")
>>> # Later...
>>> model = ContinuousTokenizer.from_pretrained("./my-tokenizer")
set_extra_state(state)[source]

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Parameters:

state (dict) – Extra state from the state_dict

Return type:

None

set_submodule(target, module, strict=False)[source]

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Parameters:

  • target (str) – The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

  • module (Module) – The module to set the submodule to.

  • strict (bool, default: False) – If False, the method will replace an existing submodule or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

  • ValueError – If the target string is empty or if module is not an instance of nn.Module.

  • AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Return type:

None

share_memory()[source]

See torch.Tensor.share_memory_().

Return type:

Self

state_dict(*args, destination=None, prefix='', keep_vars=False)[source]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

  • destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']
to(*args, **kwargs)[source]

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)[source]
to(dtype, non_blocking=False)[source]
to(tensor, non_blocking=False)[source]
to(memory_format=torch.channels_last)[source]

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

  • device (torch.device) – the desired device of the parameters and buffers in this module

  • dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

  • tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

  • memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
to_empty(*, device, recurse=True)[source]

Move the parameters and buffers to the specified device without copying storage.

Parameters:

  • device (torch.device) – The desired device of the parameters and buffers in this module.

  • recurse (bool) – Whether parameters and buffers of submodules should be recursively moved to the specified device.

Returns:

self

Return type:

Module

train(mode=True)[source]

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Parameters:

dst_type (type or string) – the desired type

Returns:

self

Return type:

Module

weights_name: str = 'pytorch_model.bin'
xpu(device=None)[source]

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

zero_grad(set_to_none=True)[source]

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Parameters:

set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.

Return type:

None

training: bool
encode(x)[source][source]

Encode input to latent representation.

Parameters:

x (Tensor | dict[str, Tensor]) – Input image/volume tensor

Return type:

tuple[Tensor, tuple[int, ...]]

Returns:

Tuple containing latent representation and additional outputs (varies by subclass - e.g., KL divergence for VAE)

decode(z, grid_shape=None)[source][source]

Decode latent representation to output.

Parameters:

  • z (Tensor) – Latent tensor (continuous codes or quantized codes)

  • grid_shape (tuple[int, ...] | None, default: None)

Return type:

Tensor

Returns:

Reconstructed output with same spatial shape as original input

forward(input)[source][source]

Full forward pass: encode -> decode.

Follows the same return contract as the sibling tokenizers (ContinuousTokenizer/DiscreteTokenizer): a dict during training (for loss computation) and a NetworkEval namedtuple during evaluation.

Parameters:

input (Tensor | dict[str, Tensor]) – Input tensor (or encoder batch dict for NeuroVFM).

Returns:

  • reconstructions: Decoded output.

    • latent/latents: Latent grid tensor.

Eval mode (NetworkEval):

  • reconstructions: Decoded output.

  • posteriors: Always None (RAE has no posterior).

  • latent: Latent grid tensor.

Return type:

Training mode (dict)

tokenize(x)[source][source]

Encode input to latent space (convenience method).

Subclasses should override with appropriate return type.

Parameters:

x (Tensor | dict[str, Tensor]) – Input tensor

Return type:

Tensor

Returns:

Latent representation (continuous or discrete indices)