vaezip 0.1.0__py3-none-any.whl

vaezip/codecs/config.py

@@ -0,0 +1,55 @@
+"""Resolving codec configs from dicts, files and shortcut names."""
+
+import copy
+import json
+import os
+from typing import Any
+
+from vaezip.codecs.shortcuts import (
+    CODEC_CLASS_KEY,
+    SHORTCUT_CONFIGS,
+    VERSION_KEY,
+    list_shortcuts,
+)
+
+ConfigSource = str | os.PathLike[str] | dict[str, Any]
+
+
+def load_config(source: ConfigSource) -> dict[str, Any]:
+    """Config dict from a config dict, a shortcut name or a path to a JSON file."""
+    if isinstance(source, dict):
+        return copy.deepcopy(source)
+    if isinstance(source, str) and source in SHORTCUT_CONFIGS:
+        return copy.deepcopy(SHORTCUT_CONFIGS[source])
+    path = os.fspath(source)
+    if not os.path.isfile(path):
+        raise ValueError(
+            f"{path!r} is neither a codec config file nor a shortcut name, "
+            f"available shortcuts: {list_shortcuts()}"
+        )
+    with open(path, encoding="utf-8") as config_file:
+        config = json.load(config_file)
+    if not isinstance(config, dict):
+        raise ValueError(f"Codec config {path!r} must contain a JSON object")
+    return config
+
+
+def split_config(
+    source: ConfigSource, overrides: dict[str, Any]
+) -> tuple[str | None, dict[str, Any]]:
+    """
+    returns:
+        class name from the config (if present), constructor parameters with the
+        overrides applied
+    """
+    config = load_config(source)
+    class_name = config.pop(CODEC_CLASS_KEY, None)
+    parameters = {
+        name: value for name, value in config.items() if not name.startswith("_")
+    }
+    parameters.update(overrides)
+    return class_name, parameters
+
+
+def without_version(config: dict[str, Any]) -> dict[str, Any]:
+    return {name: value for name, value in config.items() if name != VERSION_KEY}

vaezip/codecs/lloyd_block_range_quant.py

@@ -0,0 +1,112 @@
+"""Blockwise range quantization with Lloyd-Max level tables."""
+
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+
+from vaezip.codecs.block_range_quant import BlockRangeQuantCodec
+from vaezip.utils.lloyd import fit_lloyd_table, fit_lloyd_tables_fast, nearest_level
+
+TABLE_KEY = "table"
+
+
+class LloydBlockRangeQuantCodec(BlockRangeQuantCodec):
+    """
+    `BlockRangeQuantCodec` where the levels inside each block range are not uniform but
+    follow a Lloyd-Max table fitted to the values of each sample and stored in the
+    encoded tensor.
+
+    lloyd_fast: fit the tables with a fast histogram-based approximation.
+    cache_table: fit one table on the first encoded batch and reuse it for all next
+    calls. This speeds up encoding only, decoding always uses the table stored in the
+    encoded tensor.
+    """
+
+    def __init__(
+        self,
+        value_bits: int = 4,
+        range_bits: int = 4,
+        channel_dim: int = 1,
+        block_size: Sequence[int] = (4, 4),
+        lloyd_fast: bool = True,
+        cache_table: bool = False,
+    ) -> None:
+        super().__init__(
+            value_bits=value_bits,
+            range_bits=range_bits,
+            channel_dim=channel_dim,
+            block_size=block_size,
+        )
+        self.lloyd_fast = lloyd_fast
+        self.cache_table = cache_table
+        self._cached_table: torch.Tensor | None = None
+
+    def flush_cache(self) -> None:
+        """Drops the cached table, the next encoded batch fits a new one."""
+        self._cached_table = None
+
+    def _quantize_normalized(
+        self, normalized: torch.Tensor, span: torch.Tensor
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """
+        args:
+            normalized: [B, C, G, V], values: [0...1], position inside the block range
+            span: [B, C, G, 1], width of the block range
+        returns:
+            codes: [B, C, G, V], dtype: torch.uint8, values: [0...2^value_bits - 1]
+            tensors: table: [B, 2^value_bits], dtype: torch.float16, sorted, [0...1]
+        """
+        batch_size = normalized.shape[0]
+        values = normalized.reshape(batch_size, -1)
+        if self._cached_table is not None:
+            tables = self._cached_table.to(values.device).unsqueeze(0)
+        else:
+            weights = (span.expand_as(normalized) ** 2).reshape(batch_size, -1)
+            if self.cache_table:
+                tables = self._fit_tables(values.reshape(1, -1), weights.reshape(1, -1))
+                self._cached_table = tables[0].clone()
+            else:
+                tables = self._fit_tables(values, weights)
+        tables = tables.expand(batch_size, -1).contiguous()
+        codes = nearest_level(values, tables).reshape(normalized.shape)
+        return codes, {TABLE_KEY: tables.to(torch.float16)}
+
+    def _dequantize_normalized(
+        self, codes: torch.Tensor, tensors: dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        args:
+            codes: [B, C, G, V], dtype: torch.uint8
+            tensors: output of `_encode_tensors()`, table: [B, 2^value_bits]
+        returns:
+            normalized: [B, C, G, V], dtype: torch.float32, values: [0...1]
+        """
+        batch_size = codes.shape[0]
+        table = tensors[TABLE_KEY].float()
+        levels = table.gather(1, codes.reshape(batch_size, -1).long())
+        return levels.reshape(codes.shape)
+
+    def _fit_tables(self, values: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+        """
+        args:
+            values, weights: [R, N]
+        returns:
+            tables: [R, 2^value_bits], dtype: torch.float32
+        """
+        levels = 2**self.value_bits
+        if self.lloyd_fast:
+            return fit_lloyd_tables_fast(values, weights, levels)
+        return torch.stack(
+            [
+                fit_lloyd_table(row_values, row_weights, levels)
+                for row_values, row_weights in zip(values, weights, strict=True)
+            ]
+        )
+
+    def get_config_parameters(self) -> dict[str, Any]:
+        return {
+            **super().get_config_parameters(),
+            "lloyd_fast": self.lloyd_fast,
+            "cache_table": self.cache_table,
+        }

vaezip/codecs/range_quant.py

@@ -0,0 +1,88 @@
+"""Per-channel range quantization to low-bit integers."""
+
+import math
+from typing import Any
+
+import torch
+
+from vaezip.codecs.base import Codec
+from vaezip.utils.bit_packing import check_bits, max_code, pack_bits, unpack_bits
+from vaezip.utils.layout import (
+    channel_from_second,
+    channel_to_second,
+    check_channel_dim,
+    shape_with_channel_second,
+)
+
+CODES_KEY = "codes"
+CHANNEL_MIN_KEY = "channel_min"
+CHANNEL_MAX_KEY = "channel_max"
+SHAPE_KEY = "shape"
+MIN_RANGE = torch.finfo(torch.float32).tiny
+
+
+class RangeQuantCodec(Codec):
+    """
+    Normalizes each channel of each sample by its min and max, then quantizes to
+    `bits`-bit integers packed into bytes. The first dimension is the batch, any number
+    of spatial dimensions is supported.
+    """
+
+    def __init__(self, bits: int = 8, channel_dim: int = 1) -> None:
+        super().__init__()
+        check_bits(bits, "bits")
+        check_channel_dim(channel_dim)
+        self.bits = bits
+        self.channel_dim = channel_dim
+
+    def _encode_tensors(self, tensor: torch.Tensor) -> dict[str, torch.Tensor]:
+        """
+        args:
+            tensor: [B, ...], floating point, one channel dimension at channel_dim
+        returns:
+            tensors: codes: [B, ceil(N * bits / 8)], channel_min, channel_max: [B, C],
+            shape: [ndim], dtype: torch.int64, N: number of values per sample
+        """
+        channel_first = channel_to_second(tensor.float(), self.channel_dim)
+        batch_size, channels = channel_first.shape[:2]
+        values = channel_first.reshape(batch_size, channels, -1)  # [B, C, S]
+        channel_min = values.amin(dim=2, keepdim=True)
+        channel_max = values.amax(dim=2, keepdim=True)
+        step = (channel_max - channel_min).clamp(min=MIN_RANGE) / max_code(self.bits)
+        codes = torch.round((values - channel_min) / step)
+        codes = codes.clamp(0, max_code(self.bits)).to(torch.uint8)
+        return {
+            CODES_KEY: pack_bits(codes.reshape(batch_size, -1), self.bits),
+            CHANNEL_MIN_KEY: channel_min.reshape(batch_size, channels),
+            CHANNEL_MAX_KEY: channel_max.reshape(batch_size, channels),
+            SHAPE_KEY: torch.tensor(
+                tensor.shape, dtype=torch.int64, device=tensor.device
+            ),
+        }
+
+    def _decode_tensors(
+        self, tensors: dict[str, torch.Tensor], dtype: torch.dtype
+    ) -> torch.Tensor:
+        """
+        args:
+            tensors: output of `_encode_tensors()`
+        returns:
+            tensor: [B, ...], dtype: dtype
+        """
+        shape = [int(size) for size in tensors[SHAPE_KEY].tolist()]
+        batch_size, channels, *spatial_shape = shape_with_channel_second(
+            shape, self.channel_dim
+        )
+        count = channels * math.prod(spatial_shape)
+        codes = unpack_bits(tensors[CODES_KEY], count, self.bits)
+        codes = codes.reshape(batch_size, channels, -1).float()  # [B, C, S]
+        channel_min = tensors[CHANNEL_MIN_KEY].reshape(batch_size, channels, 1)
+        channel_max = tensors[CHANNEL_MAX_KEY].reshape(batch_size, channels, 1)
+        step = (channel_max - channel_min).clamp(min=MIN_RANGE) / max_code(self.bits)
+        values = (codes * step + channel_min).reshape(
+            batch_size, channels, *spatial_shape
+        )
+        return channel_from_second(values, self.channel_dim).to(dtype).contiguous()
+
+    def get_config_parameters(self) -> dict[str, Any]:
+        return {"bits": self.bits, "channel_dim": self.channel_dim}

vaezip/codecs/shortcuts.py

@@ -0,0 +1,49 @@
+"""Predefined codec configs addressed by short names."""
+
+from typing import Any
+
+CODEC_CLASS_KEY = "_class_name"
+VERSION_KEY = "_vaezip_version"
+DEFAULT_SHORTCUT = "cast_fp32"
+
+BLOCK_QUANT_BITS = (8, 6, 5, 4, 3, 2)
+RANGE_QUANT_BITS = (8, 6, 4)
+BLOCK_QUANT_RANGE_BITS = 4
+BLOCK_QUANT_BLOCK_EXTENT = 8
+
+
+def _build_shortcut_configs() -> dict[str, dict[str, Any]]:
+    configs: dict[str, dict[str, Any]] = {
+        DEFAULT_SHORTCUT: {CODEC_CLASS_KEY: "CastCodec", "cast_dtype": "float32"},
+        "cast_fp16": {CODEC_CLASS_KEY: "CastCodec", "cast_dtype": "float16"},
+        "cast_bf16": {CODEC_CLASS_KEY: "CastCodec", "cast_dtype": "bfloat16"},
+        "cast_fp8_e4m3": {CODEC_CLASS_KEY: "CastCodec", "cast_dtype": "float8_e4m3fn"},
+    }
+    for bits in RANGE_QUANT_BITS:
+        configs[f"rq_{bits}bit"] = {
+            CODEC_CLASS_KEY: "RangeQuantCodec",
+            "bits": bits,
+            "channel_dim": 1,
+        }
+    block_size = [BLOCK_QUANT_BLOCK_EXTENT, BLOCK_QUANT_BLOCK_EXTENT]
+    for prefix, class_name in (
+        ("brq", "BlockRangeQuantCodec"),
+        ("lbrq", "LloydBlockRangeQuantCodec"),
+    ):
+        for bits in BLOCK_QUANT_BITS:
+            configs[f"{prefix}_{bits}bit"] = {
+                CODEC_CLASS_KEY: class_name,
+                "value_bits": bits,
+                "range_bits": BLOCK_QUANT_RANGE_BITS,
+                "channel_dim": 1,
+                "block_size": list(block_size),
+            }
+    return configs
+
+
+SHORTCUT_CONFIGS = _build_shortcut_configs()
+
+
+def list_shortcuts() -> list[str]:
+    """Names of all predefined codec configs."""
+    return list(SHORTCUT_CONFIGS)

vaezip/encoded_tensor.py

@@ -0,0 +1,132 @@
+"""Container for compressed latents and its (de)serialization."""
+
+from __future__ import annotations
+
+import io
+import json
+import struct
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+from safetensors.torch import load as load_safetensors
+from safetensors.torch import save as save_safetensors
+
+CODEC_CONFIG_KEY = "codec_config"
+SAFETENSORS_FORMAT = "safetensors"
+PTH_FORMAT = "pth"
+SUPPORTED_FORMATS = (SAFETENSORS_FORMAT, PTH_FORMAT)
+
+_PTH_TENSORS_KEY = "tensors"
+_PTH_METADATA_KEY = "metadata"
+_SAFETENSORS_HEADER_SIZE_BYTES = 8
+_SAFETENSORS_METADATA_KEY = "__metadata__"
+
+
+def check_format(format: str) -> None:
+    if format not in SUPPORTED_FORMATS:
+        raise ValueError(
+            f"Unsupported format {format!r}, expected one of {SUPPORTED_FORMATS}"
+        )
+
+
+@dataclass
+class EncodedTensor:
+    """
+    Output of `Codec.encode()` and input of `Codec.decode()`.
+
+    tensors: named tensors with the compressed representation.
+    metadata: optional string metadata, normally holds the JSON config of the codec.
+    """
+
+    tensors: dict[str, torch.Tensor]
+    metadata: dict[str, str] = field(default_factory=dict)
+
+    @property
+    def codec_config(self) -> dict[str, Any] | None:
+        """Parsed codec config from the metadata, `None` if there is none."""
+        raw_config = self.metadata.get(CODEC_CONFIG_KEY)
+        if raw_config is None:
+            return None
+        config = json.loads(raw_config)
+        if not isinstance(config, dict):
+            raise ValueError("Codec config in the metadata must be a JSON object")
+        return config
+
+    @property
+    def nbytes(self) -> int:
+        """Total size of the stored tensors in bytes."""
+        return sum(
+            tensor.numel() * tensor.element_size() for tensor in self.tensors.values()
+        )
+
+    def to(self, device: torch.device | str) -> EncodedTensor:
+        return EncodedTensor(
+            tensors={name: tensor.to(device) for name, tensor in self.tensors.items()},
+            metadata=dict(self.metadata),
+        )
+
+    def serialize(
+        self, save_metadata: bool = True, format: str = SAFETENSORS_FORMAT
+    ) -> bytes:
+        """
+        args:
+            save_metadata: store the metadata (including the codec config)
+            format: "safetensors" or "pth"
+        returns:
+            serialized bytes
+        """
+        check_format(format)
+        metadata = dict(self.metadata) if save_metadata else {}
+        cpu_tensors = {
+            name: tensor.detach().to("cpu", copy=True).contiguous()
+            for name, tensor in self.tensors.items()
+        }
+        if format == SAFETENSORS_FORMAT:
+            return save_safetensors(cpu_tensors, metadata=metadata or None)
+        buffer = io.BytesIO()
+        torch.save({_PTH_TENSORS_KEY: cpu_tensors, _PTH_METADATA_KEY: metadata}, buffer)
+        return buffer.getvalue()
+
+    @classmethod
+    def deserialize(
+        cls, data: bytes, format: str = SAFETENSORS_FORMAT
+    ) -> EncodedTensor:
+        """
+        args:
+            data: bytes produced by `serialize()`
+            format: "safetensors" or "pth"
+        returns:
+            encoded tensor with tensors on CPU
+        """
+        check_format(format)
+        if format == SAFETENSORS_FORMAT:
+            return cls(
+                tensors=load_safetensors(data),
+                metadata=_read_safetensors_metadata(data),
+            )
+        content = torch.load(io.BytesIO(data), map_location="cpu", weights_only=True)
+        return cls(
+            tensors=dict(content[_PTH_TENSORS_KEY]),
+            metadata=dict(content[_PTH_METADATA_KEY]),
+        )
+
+
+def serialize(
+    encoded: EncodedTensor,
+    save_metadata: bool = True,
+    format: str = SAFETENSORS_FORMAT,
+) -> bytes:
+    return encoded.serialize(save_metadata=save_metadata, format=format)
+
+
+def deserialize(data: bytes, format: str = SAFETENSORS_FORMAT) -> EncodedTensor:
+    return EncodedTensor.deserialize(data, format=format)
+
+
+def _read_safetensors_metadata(data: bytes) -> dict[str, str]:
+    (header_size,) = struct.unpack_from("<Q", data, 0)
+    header_end = _SAFETENSORS_HEADER_SIZE_BYTES + header_size
+    header = json.loads(bytes(data[_SAFETENSORS_HEADER_SIZE_BYTES:header_end]))
+    metadata: dict[str, str] = header.get(_SAFETENSORS_METADATA_KEY, {})
+    return metadata

vaezip/utils/__init__.py

@@ -0,0 +1 @@
+"""Tensor helpers shared by the codecs."""

vaezip/utils/bit_packing.py

@@ -0,0 +1,58 @@
+"""Packing of low-bit integer codes into bytes."""
+
+import torch
+import torch.nn.functional as functional
+
+BITS_PER_BYTE = 8
+
+
+def check_bits(bits: int, name: str = "bits") -> None:
+    if isinstance(bits, bool) or not isinstance(bits, int):
+        raise TypeError(f"{name} must be an int, got {type(bits).__name__}")
+    if not 1 <= bits <= BITS_PER_BYTE:
+        raise ValueError(f"{name} must be in [1, {BITS_PER_BYTE}], got {bits}")
+
+
+def max_code(bits: int) -> int:
+    return (1 << bits) - 1
+
+
+def pack_bits(codes: torch.Tensor, bits: int) -> torch.Tensor:
+    """
+    args:
+        codes: [B, N], dtype: torch.uint8, values: [0...2^bits - 1]
+        bits: number of bits per code, in [1...8]
+    returns:
+        packed: [B, ceil(N * bits / 8)], dtype: torch.uint8, codes stored LSB first
+    """
+    if bits == BITS_PER_BYTE:
+        return codes.contiguous()
+    batch_size = codes.shape[0]
+    bit_shifts = torch.arange(bits, device=codes.device, dtype=torch.uint8)
+    bit_planes = ((codes.unsqueeze(-1) >> bit_shifts) & 1).reshape(batch_size, -1)
+    bit_planes = functional.pad(bit_planes, (0, (-bit_planes.shape[1]) % BITS_PER_BYTE))
+    byte_weights = 1 << torch.arange(
+        BITS_PER_BYTE, device=codes.device, dtype=torch.uint8
+    )
+    byte_groups = bit_planes.reshape(batch_size, -1, BITS_PER_BYTE) * byte_weights
+    return byte_groups.sum(dim=-1, dtype=torch.uint8)
+
+
+def unpack_bits(packed: torch.Tensor, count: int, bits: int) -> torch.Tensor:
+    """
+    args:
+        packed: [B, ceil(count * bits / 8)], dtype: torch.uint8
+        count: number of codes per sample
+        bits: number of bits per code, in [1...8]
+    returns:
+        codes: [B, count], dtype: torch.uint8, values: [0...2^bits - 1]
+    """
+    if bits == BITS_PER_BYTE:
+        return packed[:, :count]
+    batch_size = packed.shape[0]
+    bit_shifts = torch.arange(BITS_PER_BYTE, device=packed.device, dtype=torch.uint8)
+    bit_planes = ((packed.unsqueeze(-1) >> bit_shifts) & 1).reshape(batch_size, -1)
+    bit_planes = bit_planes[:, : count * bits]
+    code_weights = 1 << torch.arange(bits, device=packed.device, dtype=torch.uint8)
+    code_groups = bit_planes.reshape(batch_size, count, bits) * code_weights
+    return code_groups.sum(dim=-1, dtype=torch.uint8)

vaezip/utils/blocks.py

@@ -0,0 +1,91 @@
+"""Splitting spatial dimensions into blocks and merging them back."""
+
+import math
+
+import torch
+
+
+def padded_size(size: int, block_extent: int) -> int:
+    return math.ceil(size / block_extent) * block_extent
+
+
+def pad_to_block_multiple(
+    tensor: torch.Tensor, block_size: tuple[int, ...]
+) -> torch.Tensor:
+    """
+    args:
+        tensor: [B, C, *spatial]
+        block_size: block extent per spatial dimension
+    returns:
+        tensor: [B, C, *padded_spatial], spatial sizes divisible by the block size,
+        padding repeats edge values so block ranges are not affected
+    """
+    for offset, block_extent in enumerate(block_size):
+        dim = 2 + offset
+        size = tensor.shape[dim]
+        target_size = padded_size(size, block_extent)
+        if target_size != size:
+            source_index = torch.arange(target_size, device=tensor.device)
+            tensor = tensor.index_select(dim, source_index.clamp(max=size - 1))
+    return tensor
+
+
+def split_into_blocks(
+    tensor: torch.Tensor, block_size: tuple[int, ...]
+) -> torch.Tensor:
+    """
+    args:
+        tensor: [B, C, *spatial]
+        block_size: block extent per spatial dimension
+    returns:
+        blocks: [B, C, G, V], G: number of blocks per channel, V: values per block
+    """
+    padded = pad_to_block_multiple(tensor, block_size)
+    batch_size, channels = padded.shape[:2]
+    block_counts = [
+        size // extent
+        for size, extent in zip(padded.shape[2:], block_size, strict=True)
+    ]
+    split_shape = [
+        part
+        for count, extent in zip(block_counts, block_size, strict=True)
+        for part in (count, extent)
+    ]
+    ndim = len(block_size)
+    count_dims = range(2, 2 + 2 * ndim, 2)
+    extent_dims = range(3, 3 + 2 * ndim, 2)
+    split = padded.reshape(batch_size, channels, *split_shape)  # [B, C, n1, b1, n2, b2]
+    grouped = split.permute(0, 1, *count_dims, *extent_dims)  # [B, C, n1, n2, b1, b2]
+    return grouped.reshape(
+        batch_size, channels, math.prod(block_counts), math.prod(block_size)
+    )
+
+
+def merge_blocks(
+    blocks: torch.Tensor, spatial_shape: list[int], block_size: tuple[int, ...]
+) -> torch.Tensor:
+    """
+    args:
+        blocks: [B, C, G, V]
+        spatial_shape: original spatial sizes before padding
+        block_size: block extent per spatial dimension
+    returns:
+        tensor: [B, C, *spatial_shape]
+    """
+    batch_size, channels = blocks.shape[:2]
+    block_counts = [
+        padded_size(size, extent) // extent
+        for size, extent in zip(spatial_shape, block_size, strict=True)
+    ]
+    ndim = len(block_size)
+    grouped = blocks.reshape(batch_size, channels, *block_counts, *block_size)
+    split_order = [
+        dim for index in range(ndim) for dim in (2 + index, 2 + ndim + index)
+    ]
+    split = grouped.permute(0, 1, *split_order)  # [B, C, n1, b1, n2, b2]
+    padded_shape = [
+        count * extent for count, extent in zip(block_counts, block_size, strict=True)
+    ]
+    padded = split.reshape(batch_size, channels, *padded_shape)
+    crop = tuple(slice(0, size) for size in spatial_shape)
+    return padded[(slice(None), slice(None), *crop)]

vaezip/utils/layout.py

@@ -0,0 +1,50 @@
+"""Tensor layout handling: batch, channel and spatial dimensions."""
+
+import torch
+
+BATCH_DIM = 0
+
+
+def check_channel_dim(channel_dim: int) -> None:
+    if isinstance(channel_dim, bool) or not isinstance(channel_dim, int):
+        raise TypeError(f"channel_dim must be an int, got {type(channel_dim).__name__}")
+    if channel_dim == BATCH_DIM:
+        raise ValueError("channel_dim can not be 0, the first dimension is the batch")
+
+
+def resolve_channel_dim(channel_dim: int, ndim: int) -> int:
+    """Non-negative channel index, negative values count from the end."""
+    if not -ndim < channel_dim < ndim or channel_dim % ndim == BATCH_DIM:
+        raise ValueError(
+            f"channel_dim={channel_dim} is invalid for a tensor with {ndim} dimensions"
+        )
+    return channel_dim % ndim
+
+
+def channel_to_second(tensor: torch.Tensor, channel_dim: int) -> torch.Tensor:
+    """
+    args:
+        tensor: [B, ...] with the channel at `channel_dim`
+    returns:
+        tensor: [B, C, ...] view with spatial dimensions in the original order
+    """
+    return tensor.movedim(resolve_channel_dim(channel_dim, tensor.ndim), 1)
+
+
+def channel_from_second(tensor: torch.Tensor, channel_dim: int) -> torch.Tensor:
+    """
+    args:
+        tensor: [B, C, ...]
+    returns:
+        tensor: [B, ...] view with the channel moved to `channel_dim`
+    """
+    return tensor.movedim(1, resolve_channel_dim(channel_dim, tensor.ndim))
+
+
+def shape_with_channel_second(shape: list[int], channel_dim: int) -> list[int]:
+    """Shape [B, C, ...] of a tensor with the given original shape."""
+    channel_index = resolve_channel_dim(channel_dim, len(shape))
+    spatial_sizes = [
+        size for index, size in enumerate(shape) if index not in (0, channel_index)
+    ]
+    return [shape[0], shape[channel_index], *spatial_sizes]