Torch Einops Kit

https://github.com/hunterhogan/torch_einops_kit
Admin
Torch Einops Kit is a typed PyTorch utility library providing tensor-shaping, masking, padding,...
Tokens:26,125
Snippets:140
Trust Score:8.3
Update:6 months ago
Show doc for...
Tokens:
Raw
### Install torch_einops_kit

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Install the library using pip or uv.

```bash
pip install torch_einops_kit
```

```bash
# or
uv add torch_einops_kit
```

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

### Install Dependencies

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Installs the necessary dependencies for the project using uv.

```bash
uv sync
```

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

### Install torch_einops_kit with pip

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Use pip to install the torch-einops-kit package.

```bash
pip install torch_einops_kit
```

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

### Install torch_einops_kit with uv

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Use uv to add the torch-einops-kit package to your project dependencies.

```bash
uv add torch_einops_kit
```

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

### Pad Left/Right at Dimension

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Prepend or append a fixed number of fill elements along a specified dimension. Useful for adding start tokens or padding to a fixed size.

```python
import torch
from torch_einops_kit import pad_left_at_dim, pad_right_at_dim

t = torch.randn(2, 5)

left_padded  = pad_left_at_dim(t, 2, dim=1)    # (2, 7)
right_padded = pad_right_at_dim(t, 2, dim=1)   # (2, 7)
assert left_padded.shape  == (2, 7)
assert right_padded.shape == (2, 7)

# Prepend a start token (pi-zero-pytorch pattern)
discrete_ids = torch.randint(0, 100, (4, 10))
with_start = pad_left_at_dim(discrete_ids + 1, 1)
assert with_start.shape == (4, 11)
```

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

### Execute a callable at most once

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Ensures a function is called only on its first invocation. Subsequent calls return `None` silently. Useful for one-time setup or logging.

```python
from torch_einops_kit import once

print_once = once(print)

for _ in range(5):
    print_once("A100 GPU detected, using flash attention")
# Prints exactly once; subsequent calls silently return None
```

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

### Import Lightweight NN Module Adapters

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Imports lightweight adapters for PyTorch's nn.Module. These include Identity, Lambda, and Sequential for building simple network structures.

```python
from torch_einops_kit.nn import (
    Identity,
    Lambda,
    Sequential,
)
```

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

### Run Static Analysis

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Performs static code analysis using ruff check.

```bash
ruff check .
```

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

### Import Device Decorators

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Imports decorators for managing tensor device placement. Use these to ensure modules and their inputs are on the correct device.

```python
from torch_einops_kit.device import (
    module_device,
    move_inputs_to_device,
    move_inputs_to_module_device,
)
```

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

### Import Normalization and Masked Reduction Helpers

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Imports utilities for normalization (RMSNorm, l2norm) and masked reductions (masked_mean, exclusive_cumsum). Useful for various scaling and aggregation tasks.

```python
from torch_einops_kit.scaleValues import (
    exclusive_cumsum,
    l2norm,
    RMSNorm,
    masked_mean,
)
```

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

### Run Test Suite

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Executes the project's test suite using pytest.

```bash
pytest
```

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

### Device Utilities

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Utilities for device inference and automatic tensor movement within PyTorch modules.

```APIDOC
## `module_device(m)`

### Description
Returns the device of the first parameter or registered buffer in `m`. Returns `None` when `m` has neither parameters nor buffers.

### Contract
`module_device(m)`
```

```APIDOC
## `move_inputs_to_device(device)`

### Description
Decorator that recursively moves every tensor inside positional and keyword arguments to `device` before calling the wrapped function. Non-tensor values pass through unchanged.

### Contract
`move_inputs_to_device(device)`
```

```APIDOC
## `move_inputs_to_module_device(fn)`

### Description
Decorator for methods whose first argument is a `torch.nn.Module`. The decorator infers the target device with `module_device(self)` and moves every tensor argument after `self` to that device. If `module_device(self)` returns `None`, the call is a no-op.

### Contract
`move_inputs_to_module_device(fn)`
```

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

### Save/Load Utilities

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Utilities for saving and loading PyTorch models, including a checkpoint decorator and configuration helpers.

```APIDOC
## `save_load(...)`

### Description
Class decorator for `torch.nn.Module` subclasses. The decorator records constructor arguments on each instance and adds instance save / load methods plus a classmethod that reconstructs a new instance from a checkpoint.

### Contract
`save_load(...)`

### Default Methods Added
- `save(path, overwrite=True)`
- `load(path, strict=True)`
- `init_and_load(path, strict=True)`
```

```APIDOC
## `dehydrate_config(config, config_instance_var_name)`

### Description
Walks nested `list`, `tuple`, and `dict` structures and replaces decorated module instances with reconstruction records.

### Contract
`dehydrate_config(config, config_instance_var_name)`
```

```APIDOC
## `rehydrate_config(config)`

### Description
Walks nested `list`, `tuple`, and `dict` structures and reconstructs decorated modules from the stored reconstruction records.

### Contract
`rehydrate_config(config)`
```

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

### Import Tensor Helpers from Package Root

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Imports common tensor utility functions from the main package namespace. Ensure these functions are available for direct use.

```python
from torch_einops_kit import (
    align_dims_left,
    and_masks,
    broadcast_cat,
    lens_to_mask,
    maybe,
    once,
    or_masks,
    pad_sequence_and_cat,
    pad_sequence,
    safe_cat,
    safe_stack,
    shape_with_replace,
    slice_at_dim,
    tree_flatten_with_inverse,
    tree_map_tensor,
)
```

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

### Neural Network Modules

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Lightweight `torch.nn.Module` adapters and a `Sequential` constructor.

```APIDOC
## `Sequential(*modules)`

### Description
Equivalent to `nn.Sequential(*compact(modules))`. Every `None` argument is discarded before construction.

### Contract
`Sequential(*modules)`
```

```APIDOC
## `Identity()`

### Description
`torch.nn.Module` wrapper around `identity`. Unlike `torch.nn.Identity`, `Identity.forward` accepts additional positional arguments and keyword arguments and returns the first positional argument unchanged.

### Contract
`Identity()`
```

```APIDOC
## `Lambda(fn)`

### Description
`torch.nn.Module` wrapper that stores `fn` on `self.fn` and delegates `forward(*args, **kwargs)` to `fn(*args, **kwargs)` without changing the argument structure.

### Contract
`Lambda(fn)`
```

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

### Import Einops Pack/Unpack Helpers

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Imports functions for packing and unpacking tensors using einops patterns. These are useful for handling variable dimensions.

```python
from torch_einops_kit.einops import (
    pack_one,
    pack_with_inverse,
    unpack_one,
)
```

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

### module_device(m)

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Infer a module's device. Returns the torch.device of the first parameter or registered buffer in m. Returns None for stateless modules.

```APIDOC
## module_device(m)

### Description
Infers and returns the `torch.device` of a given `torch.nn.Module`. It checks the device of the module's parameters or registered buffers. Returns `None` if the module is stateless (i.e., has no parameters or buffers).

### Parameters
* **m** (torch.nn.Module) - The module whose device needs to be inferred.

### Returns
* **device** (torch.device or None) - The device of the module, or None if stateless.

### Usage
```python
import torch
from torch import nn
from torch_einops_kit.device import module_device

linear = nn.Linear(3, 5)
assert module_device(linear) == torch.device("cpu")

# GPU module
if torch.cuda.is_available():
    gpu_linear = nn.Linear(3, 5).cuda()
    assert module_device(gpu_linear) == torch.device("cuda", 0)

# Stateless module returns None
identity = nn.Identity()
assert module_device(identity) is None
```
```

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

### Import Checkpoint Save/Load Decorators

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Imports decorators for saving and loading model configurations and states. These helpers simplify checkpointing.

```python
from torch_einops_kit.save_load import (
    dehydrate_config,
    rehydrate_config,
    save_load,
)
```

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

### einops Submodule

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Utilities for packing and unpacking tensors using einops patterns, with paired inverse functions.

```APIDOC
## pack_one(t, pattern)

### Description
Packs one tensor using an einops pattern and returns the packed tensor and shape metadata for paired reconstruction with `unpack_one`.

### Parameters
#### Path Parameters
- **t** (tensor) - The tensor to pack.
- **pattern** (str) - The einops pattern to use for packing.

### Return
- **tuple** - A tuple containing the packed tensor and shape metadata.
```

```APIDOC
## pack_with_inverse(t, pattern)

### Description
Wraps `einops.pack` and stores the corresponding inverse. Input `t` may be one tensor or a `list[Tensor]`. The inverse returns the matching kind and optionally accepts a different `inv_pattern` for unpacking.

### Parameters
#### Path Parameters
- **t** (tensor or list[tensor]) - The tensor(s) to pack.
- **pattern** (str) - The einops pattern to use for packing.

### Return
- **tuple** - A tuple containing the packed tensor and the inverse function.
```

```APIDOC
## unpack_one(t, ps, pattern)

### Description
Unpacks one tensor using packed-shape metadata produced by `pack_one`.

### Parameters
#### Path Parameters
- **t** (tensor) - The tensor to unpack.
- **ps** (any) - The packed-shape metadata.
- **pattern** (str) - The einops pattern to use for unpacking.

### Return
- **tensor** - The unpacked tensor.
```

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

### pad_left_at_dim_to / pad_right_at_dim_to

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

These functions pad a tensor to a minimum target length along a specified dimension. If the tensor is already long enough, it is returned unchanged.

```APIDOC
## pad_left_at_dim_to(t, length, dim)

### Description
Pads a tensor to a minimum target length by prepending elements along a specified dimension.

### Parameters
#### Path Parameters
- **t** (torch.Tensor) - Input tensor.
- **length** (int) - The target minimum length.
- **dim** (int) - The dimension along which to pad.

### Request Example
```python
import torch
from torch_einops_kit import pad_left_at_dim_to

t = torch.randn(3, 6, 1)
result = pad_left_at_dim_to(t, 7, dim=1)
# result.shape == (3, 7, 1)
```

## pad_right_at_dim_to(t, length, dim)

### Description
Pads a tensor to a minimum target length by appending elements along a specified dimension.

### Parameters
#### Path Parameters
- **t** (torch.Tensor) - Input tensor.
- **length** (int) - The target minimum length.
- **dim** (int) - The dimension along which to pad.

### Request Example
```python
import torch
from torch_einops_kit import pad_right_at_dim_to

t = torch.randn(3, 6, 1)
result = pad_right_at_dim_to(t, 7, dim=1)
# result.shape == (3, 7, 1)
```
```

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

### Decorate Module with Save/Load Helpers

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Applies a decorator to an nn.Module to enable easy saving and loading of its state and configuration. The decorated class gains `save` and `init_and_load` methods.

```python
from pathlib import Path

from torch import Tensor, nn

from torch_einops_kit.save_load import save_load


@save_load()
class SimpleNet(nn.Module):
    def __init__(self, dim: int, hidden_dim: int) -> None:
        super().__init__()
        self.dim = dim
        self.hidden_dim = hidden_dim
        self.net = nn.Linear(dim, hidden_dim)

    def forward(self, x: Tensor) -> Tensor:
        return self.net(x)


path = Path("model.pt")
model = SimpleNet(10, 20)
model.save(str(path))

restored = SimpleNet.init_and_load(str(path))
assert restored.dim == 10
assert restored.hidden_dim == 20
```

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

### Insert Singleton Dimensions

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Functions `pad_ndim`, `pad_left_ndim`, and `pad_right_ndim` insert singleton dimensions to increase a tensor's rank for broadcasting. `pad_ndim` adds dimensions to both sides, while `pad_left_ndim` and `pad_right_ndim` add them to the beginning or end, respectively. No data is copied.

```python
import torch
from torch_einops_kit import pad_ndim, pad_left_ndim, pad_right_ndim

t = torch.randn(4, 8)          # shape: (4, 8)

both  = pad_ndim(t, (1, 2))    # (1, 4, 8, 1, 1)
left  = pad_left_ndim(t, 2)    # (1, 1, 4, 8)
right = pad_right_ndim(t, 2)   # (4, 8, 1, 1)

assert both.shape  == (1, 4, 8, 1, 1)
assert left.shape  == (1, 1, 4, 8)
assert right.shape == (4, 8, 1, 1)
```

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

### Pad to Target Rank

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Use `pad_left_ndim_to` or `pad_right_ndim_to` to prepend or append singleton dimensions until a tensor reaches a specified rank. If the tensor already meets or exceeds the target rank, it is returned unchanged. This is useful for broadcasting scalars or lower-rank tensors against higher-rank tensors.

```python
import torch
from torch_einops_kit import pad_right_ndim_to

# Broadcast a 1-D time scalar against a 5-D video tensor (b, c, t, h, w)
video = torch.randn(2, 3, 10, 64, 64)
time  = torch.tensor([0.5, 0.3])          # shape: (2,)

aligned_time = pad_right_ndim_to(time, video.ndim)  # (2, 1, 1, 1, 1)
noised = video * aligned_time                        # broadcasts correctly
assert noised.shape == (2, 3, 10, 64, 64)
```

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

### Build Mask using exists and lens_to_mask

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Uses the `exists` TypeGuard to check for optional values and `lens_to_mask` to create a mask from tensor lengths. The `max_len` parameter can be optionally provided.

```python
import torch
from torch_einops_kit import exists, lens_to_mask

def build_mask(lens: torch.Tensor, max_len: int | None = None) -> torch.Tensor:
    if not exists(max_len):
        max_len = int(lens.amax().item())
    return lens_to_mask(lens, max_len)

lens = torch.tensor([4, 3, 1])
mask = build_mask(lens)
assert mask.shape == (3, 4)
mask_wide = build_mask(lens, max_len=6)
assert mask_wide.shape == (3, 6)
```

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

### move_inputs_to_device(device)

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Decorator for explicit device routing. Creates a decorator that moves all tensor arguments to a fixed target device before each call.

```APIDOC
## move_inputs_to_device(device)

### Description
A decorator that automatically moves all tensor arguments of a function to a specified target device before the function is called. Non-tensor arguments are passed through unchanged.

### Parameters
* **device** (torch.device) - The target device to move tensor arguments to.

### Usage
```python
import torch
from torch_einops_kit.device import move_inputs_to_device

target = torch.device("meta")

@move_inputs_to_device(target)
def compute(
    x: torch.Tensor,
    nested: tuple[torch.Tensor, str],
    *, 
    kw: torch.Tensor,
) -> tuple[torch.device, torch.device, torch.device]:
    return x.device, nested[0].device, kw.device

cpu_t = torch.tensor([1.0, 2.0])
result = compute(cpu_t, (cpu_t, "tag"), kw=cpu_t)
# All tensors moved to "meta" device; non-tensor values pass through unchanged
assert all(d == torch.device("meta") for d in result)
```
```

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

### Conditional function application with maybe

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Applies a function `fn` only if the first argument is not `None`. If `fn` is `None`, it defaults to an identity function. Returns `None` if the first argument is `None`.

```python
import torch
from torch_einops_kit import maybe, lens_to_mask

# episode_lens may be None at the start of training
episode_lens: torch.Tensor | None = torch.tensor([4, 3, 1])
seq_len = 5

# mask is None when episode_lens is None, otherwise lens_to_mask is applied
mask = maybe(lens_to_mask)(episode_lens, seq_len)
assert mask is not None
assert mask.shape == (3, 5)

# When episode_lens is None, mask is None
mask_none = maybe(lens_to_mask)(None, seq_len)
assert mask_none is None

# Pass None as fn to get identity behavior
result = maybe(None)(42)
assert result == 42
```

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

### pad_left_at_dim / pad_right_at_dim

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

These functions prepend or append a fixed number of fill elements along a specified dimension of a tensor.

```APIDOC
## pad_left_at_dim(t, pad, dim)

### Description
Prepends a fixed number of fill elements along one dimension.

### Parameters
#### Path Parameters
- **t** (torch.Tensor) - Input tensor.
- **pad** (int) - Number of elements to prepend.
- **dim** (int) - The dimension along which to pad.

### Request Example
```python
import torch
from torch_einops_kit import pad_left_at_dim

t = torch.randn(2, 5)
left_padded = pad_left_at_dim(t, 2, dim=1)
# left_padded.shape == (2, 7)
```

## pad_right_at_dim(t, pad, dim)

### Description
Appends a fixed number of fill elements along one dimension.

### Parameters
#### Path Parameters
- **t** (torch.Tensor) - Input tensor.
- **pad** (int) - Number of elements to append.
- **dim** (int) - The dimension along which to pad.

### Request Example
```python
import torch
from torch_einops_kit import pad_right_at_dim

t = torch.randn(2, 5)
right_padded = pad_right_at_dim(t, 2, dim=1)
# right_padded.shape == (2, 7)
```
```

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

### Checkpoint decorator for nn.Module subclasses

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Use the `save_load` decorator to automatically add save and load capabilities to your `nn.Module` subclasses. It records constructor arguments for easy reconstruction. Custom method names and version tagging are supported.

```python
from pathlib import Path
import torch
from torch import nn, Tensor
from torch_einops_kit.save_load import save_load

@save_load()
class SimpleNet(nn.Module):
    def __init__(self, dim: int, hidden_dim: int) -> None:
        super().__init__()
        self.dim = dim
        self.hidden_dim = hidden_dim
        self.net = nn.Linear(dim, hidden_dim)

    def forward(self, x: Tensor) -> Tensor:
        return self.net(x)

path = Path("model.pt")
model = SimpleNet(10, 20)
model.save(str(path))

# Reconstruct from checkpoint without knowing constructor args
restored = SimpleNet.init_and_load(str(path))
assert restored.dim == 10
assert restored.hidden_dim == 20
```

```python
# Custom method names and version tagging
@save_load(
    save_method_name="store",
    load_method_name="restore",
    config_instance_var_name="stored_config",
    init_and_load_classmethod_name="create_and_restore",
    version="1.0.0",
)
class VersionedNet(nn.Module):
    def __init__(self, width: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.randn(width))

net = VersionedNet(13)
net.store("versioned_model.pt")
restored_net = VersionedNet.create_and_restore("versioned_model.pt")
assert restored_net.weight.shape == (13,)
```

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

### Batch Variable-Length Tensors and Build Mask

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Pads variable-length tensors and generates a corresponding mask. Useful for batching sequences of different lengths.

```python
import torch

from torch_einops_kit import lens_to_mask, pad_sequence

x = torch.randn(2, 4, 5)
y = torch.randn(2, 3, 5)
z = torch.randn(2, 1, 5)

packed, lens = pad_sequence([x, y, z], dim=1, return_lens=True)
mask = lens_to_mask(lens)

assert packed.shape == (3, 2, 4, 5)
assert lens.tolist() == [4, 3, 1]
assert torch.allclose(mask.sum(dim=-1), lens)
```

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

### move_inputs_to_module_device(fn)

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Decorator for methods of torch.nn.Module subclasses. Infers the target device from self and moves all tensor arguments to it before each call.

```APIDOC
## move_inputs_to_module_device(fn)

### Description
A decorator designed for methods within `torch.nn.Module` subclasses. It automatically infers the target device from the module instance (`self`) and moves all tensor arguments passed to the decorated method to that device before execution. This is useful for ensuring consistency in device placement for module inputs.

### Usage
```python
import torch
from torch import nn, Tensor
from torch_einops_kit.device import move_inputs_to_module_device

class EchoModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.scale = nn.Parameter(torch.tensor([2.0], device=torch.device("meta")))

    @move_inputs_to_module_device
    def forward(self, x: Tensor) -> Tensor:
        return x

module = EchoModule()
cpu_tensor = torch.tensor([1.0, 2.0])
result = module.forward(cpu_tensor)
assert result.device == torch.device("meta")

# Standalone function attached to a module class
@move_inputs_to_module_device
def policy_loss(model: nn.Module, states: Tensor, actions: Tensor) -> Tensor:
    return (states + actions).mean()
```
```

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

### Infer Module's Device

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Returns the torch.device of the first parameter or registered buffer in a module. Returns None for stateless modules.

```python
import torch
from torch import nn
from torch_einops_kit.device import module_device

linear = nn.Linear(3, 5)
assert module_device(linear) == torch.device("cpu")

# GPU module
if torch.cuda.is_available():
    gpu_linear = nn.Linear(3, 5).cuda()
    assert module_device(gpu_linear) == torch.device("cuda", 0)

# Stateless module returns None
identity = nn.Identity()
assert module_device(identity) is None
```

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

### Decorator to Move Inputs to a Fixed Device

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Creates a decorator that moves all tensor arguments to a fixed target device before each call. Non-tensor arguments are passed through unchanged.

```python
import torch
from torch_einops_kit.device import move_inputs_to_device

target = torch.device("meta")

@move_inputs_to_device(target)
def compute(
    x: torch.Tensor,
    nested: tuple[torch.Tensor, str],
    *, 
    kw: torch.Tensor,
) -> tuple[torch.device, torch.device, torch.device]:
    return x.device, nested[0].device, kw.device

cpu_t = torch.tensor([1.0, 2.0])
result = compute(cpu_t, (cpu_t, "tag"), kw=cpu_t)
# All tensors moved to "meta" device; non-tensor values pass through unchanged
assert all(d == torch.device("meta") for d in result)
```

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

### Nested module checkpoint serialization

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Use `dehydrate_config` and `rehydrate_config` to serialize and reconstruct nested decorated module instances. `init_and_load` automatically handles rehydration of nested modules.

```python
from torch import nn
import torch
from torch_einops_kit.save_load import save_load, dehydrate_config, rehydrate_config

@save_load()
class InnerNet(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        self.proj = nn.Linear(dim, dim)

@save_load()
class OuterNet(nn.Module):
    def __init__(self, inner: InnerNet, scale: float) -> None:
        super().__init__()
        self.inner = inner
        self.scale = nn.Parameter(torch.tensor([scale]))

    def forward(self, x):
        return self.inner.proj(x) * self.scale

inner = InnerNet(64)
outer = OuterNet(inner, scale=2.0)
outer.save("outer_model.pt")

# init_and_load rehydrates the nested InnerNet automatically
restored_outer = OuterNet.init_and_load("outer_model.pt")
assert isinstance(restored_outer.inner, InnerNet)
```

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

### dehydrate_config / rehydrate_config

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Functions for handling nested module checkpoint serialization. `dehydrate_config` replaces nested decorated module instances with checkpoint-safe reconstruction records, and `rehydrate_config` reconstructs those modules from the stored records.

```APIDOC
### `dehydrate_config` / `rehydrate_config` — nested module checkpoint serialization

`dehydrate_config` replaces nested decorated module instances with checkpoint-safe reconstruction records. `rehydrate_config` reconstructs those modules from the stored records.

```python
from torch import nn
import torch
from torch_einops_kit.save_load import save_load, dehydrate_config, rehydrate_config

@save_load()
class InnerNet(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        self.proj = nn.Linear(dim, dim)

@save_load()
class OuterNet(nn.Module):
    def __init__(self, inner: InnerNet, scale: float) -> None:
        super().__init__()
        self.inner = inner
        self.scale = nn.Parameter(torch.tensor([scale]))

    def forward(self, x):
        return self.inner.proj(x) * self.scale

inner = InnerNet(64)
outer = OuterNet(inner, scale=2.0)
outer.save("outer_model.pt")

# init_and_load rehydrates the nested InnerNet automatically
restored_outer = OuterNet.init_and_load("outer_model.pt")
assert isinstance(restored_outer.inner, InnerNet)
```
```

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

### pad_left_ndim_to / pad_right_ndim_to

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Functions to pad a tensor with singleton dimensions on the left or right until it reaches a target rank. If the tensor already has at least the target rank, it is returned unchanged.

```APIDOC
## pad_left_ndim_to(t, ndims) / pad_right_ndim_to(t, ndims)

### Description
Prepend or append singleton dimensions until a tensor has at least `ndims` dimensions. Returns `t` unchanged when `t.ndim >= ndims`.

### Parameters
- **t** (torch.Tensor) - The input tensor.
- **ndims** (int) - The target minimum number of dimensions.

### Example
```python
import torch
from torch_einops_kit import pad_right_ndim_to

# Broadcast a 1-D time scalar against a 5-D video tensor (b, c, t, h, w)
video = torch.randn(2, 3, 10, 64, 64)
time  = torch.tensor([0.5, 0.3])          # shape: (2,)

aligned_time = pad_right_ndim_to(time, video.ndim)  # (2, 1, 1, 1, 1)
noised = video * aligned_time                        # broadcasts correctly
assert noised.shape == (2, 3, 10, 64, 64)
```
```

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

### pack_with_inverse(t, pattern)

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Packs a tensor (or list of tensors) using an einops pattern and returns a paired inverse function for shape restoration.

```APIDOC
## pack_with_inverse(t, pattern)

### Description
Packs a tensor or a list of tensors using an einops pattern and returns a paired inverse function for shape restoration.

### Parameters
* **t** (torch.Tensor or list[torch.Tensor]) - The tensor(s) to pack.
* **pattern** (str) - The einops pattern to use for packing.

### Returns
* **packed** (torch.Tensor) - The packed tensor.
* **inverse** (function) - A function to restore the original shape.

### Usage
```python
import torch
from torch_einops_kit.einops import pack_with_inverse

# Pack a single tensor: flatten spatial dims
t = torch.randn(3, 12, 2, 2)
packed, inverse = pack_with_inverse(t, "b * d")
assert packed.shape == (3, 24, 2)

# Recover original shape
restored = inverse(packed)
assert restored.shape == (3, 12, 2, 2)

# Pack a list of tensors and unpack with an overriding pattern
t = torch.randn(3, 12, 2)
u = torch.randn(3, 4, 2)
packed, inverse = pack_with_inverse([t, u], "b * d")
assert packed.shape == (3, 28, 2)

# Apply an operation on the packed tensor, then unpack
reduced = packed.sum(dim=-1)        # (3, 28)
t_out, u_out = inverse(reduced, "b *")
assert t_out.shape == (3, 12)
assert u_out.shape == (3, 4)
```
```

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

### Pad Left/Right to Target Length

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Pad tensors to a minimum target length along a specified dimension. The tensor is returned unchanged if it's already long enough.

```python
import torch
from torch_einops_kit import pad_left_at_dim_to, pad_right_at_dim_to

t = torch.randn(3, 6, 1)

# Pad to length 7 along dim 1
result = pad_left_at_dim_to(t, 7, dim=1)
assert result.shape == (3, 7, 1)

# Already at length 6; tensor returned unchanged
result_same = pad_left_at_dim_to(t, 6, dim=1)
assert result_same is t

# Align variable-length action sequences to max_time before batching (dreamer4)
sequences = [torch.randn(2, length, 32) for length in [4, 7, 5]]
max_time = 8
aligned = [pad_right_at_dim_to(s, max_time, dim=1) for s in sequences]
assert all(s.shape == (2, 8, 32) for s in aligned)
```

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

### Pack Tensor with Einops and Keep Inverse Function

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Packs a tensor using an einops pattern and retains a function to restore the original shape. Useful for simplifying tensor dimensions while maintaining reconstructability.

```python
import torch

from torch_einops_kit import pack_with_inverse

t = torch.randn(3, 12, 2, 2)
packed, inverse = pack_with_inverse(t, "b * d")

assert packed.shape == (3, 24, 2)

restored = inverse(packed)
assert restored.shape == (3, 12, 2, 2)
```

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

### Mask helpers

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Utility functions for creating and manipulating boolean masks.

```APIDOC
## lens_to_mask(lens, max_len=None)

### Description
Converts integer length values to a boolean mask.

### Parameters
#### Path Parameters
- **lens** (tensor) - A tensor containing integer length values.
- **max_len** (int, optional) - The maximum length for the mask. If None, it's inferred from `lens`. Defaults to None.

### Returns
A boolean mask with shape `(*lens.shape, max_len)`. Position `i` in the final axis is `True` when `i < lens[...]`.
```

```APIDOC
## reduce_masks(masks, op)

### Description
Filters out `None` values from a list of masks, then reduces the remaining masks left-to-right using the specified operation `op`.

### Parameters
#### Path Parameters
- **masks** (list of boolean tensors or None) - A list of masks to reduce.
- **op** (callable) - The reduction operation (e.g., `torch.logical_and`, `torch.logical_or`).

### Returns
- The reduced mask. Returns `None` when no non-`None` masks remain.
```

```APIDOC
## and_masks(masks)

### Description
Equivalent to `reduce_masks(masks, torch.logical_and)`. Reduces a list of masks using logical AND.

### Parameters
#### Path Parameters
- **masks** (list of boolean tensors or None) - A list of masks to reduce.

### Returns
- The combined mask using logical AND. Returns `None` when no active mask remains.
```

```APIDOC
## or_masks(masks)

### Description
Equivalent to `reduce_masks(masks, torch.logical_or)`. Reduces a list of masks using logical OR.

### Parameters
#### Path Parameters
- **masks** (list of boolean tensors or None) - A list of masks to reduce.

### Returns
- The combined mask using logical OR. Returns `None` when no active mask remains.
```

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

### pad_left_at_dim

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Prepends a specified number of values along a given dimension of a tensor.

```APIDOC
## pad_left_at_dim

### Description
Prepends `pad` values along `dim`.

### Parameters
- **t** (Tensor) - The input tensor.
- **pad** (int) - The number of values to prepend.
- **dim** (int, optional) - The dimension along which to pad. Defaults to -1.
- **value** (float, optional) - The padding value. Defaults to 0.0.
```

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

### pad_ndim / pad_left_ndim / pad_right_ndim

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Functions to insert singleton dimensions into a tensor's shape, either on both sides, the left, or the right. This is useful for preparing tensors for broadcasting operations.

```APIDOC
## pad_ndim(t, (left, right)) / pad_left_ndim(t, ndims) / pad_right_ndim(t, ndims)

### Description
Insert singleton dimensions on either side of a tensor shape to increase its rank for broadcasting. No data is copied.

### Parameters
- **t** (torch.Tensor) - The input tensor.
- **(left, right)** (tuple of int) - For `pad_ndim`, the number of singleton dimensions to add to the left and right, respectively.
- **ndims** (int) - For `pad_left_ndim` and `pad_right_ndim`, the number of singleton dimensions to add to the left or right.

### Example
```python
import torch
from torch_einops_kit import pad_ndim, pad_left_ndim, pad_right_ndim

t = torch.randn(4, 8)          # shape: (4, 8)

both  = pad_ndim(t, (1, 2))    # (1, 4, 8, 1, 1)
left  = pad_left_ndim(t, 2)    # (1, 1, 4, 8)
right = pad_right_ndim(t, 2)   # (4, 8, 1, 1)

assert both.shape  == (1, 4, 8, 1, 1)
assert left.shape  == (1, 1, 4, 8)
assert right.shape == (4, 8, 1, 1)
```
```

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

### pack_one(t, pattern) / unpack_one(t, ps, pattern)

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Pack a single tensor and store shape metadata; recover the original shape from that metadata.

```APIDOC
## pack_one(t, pattern) / unpack_one(t, ps, pattern)

### Description
Pack a single tensor and store shape metadata; recover the original shape from that metadata.

### Parameters
* **t** (torch.Tensor) - The tensor to pack.
* **pattern** (str) - The einops pattern to use for packing.
* **ps** (tuple) - The shape metadata returned by `pack_one`.

### Returns
* **packed** (torch.Tensor) - The packed tensor.
* **ps** (tuple) - Shape metadata for unpacking.
* **restored** (torch.Tensor) - The tensor with its original shape restored.

### Usage
```python
import torch
from torch_einops_kit.einops import pack_one, unpack_one

t = torch.randn(4, 8, 3, 3)
packed, ps = pack_one(t, "b * d")
assert packed.shape == (4, 24, 3)

restored = unpack_one(packed, ps, "b * d")
assert restored.shape == (4, 8, 3, 3)
```
```

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

### Concatenation and stacking helpers

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Functions for safely concatenating and stacking tensors, handling None values and broadcasting.

```APIDOC
## broadcast_cat(tensors, dim=0)

### Description
Broadcasts tensor groups before concatenation along the specified dimension.

### Parameters
#### Path Parameters
- **tensors** (list of tensors) - A list of tensors to concatenate.
- **dim** (int, optional) - The dimension along which to concatenate. Defaults to 0.
```

```APIDOC
## safe_stack(tensors, dim=0)

### Description
Removes `None` values from a list of tensors, then applies `torch.stack` along the specified dimension. Even a single tensor is stacked, so the result gains a new dimension.

### Parameters
#### Path Parameters
- **tensors** (list of tensors or None) - A list of tensors to stack.
- **dim** (int, optional) - The dimension along which to stack. Defaults to 0.

### Returns
- The stacked tensor. Returns `None` when no active tensor remains.
```

```APIDOC
## safe_cat(tensors, dim=0)

### Description
Removes `None` values from a list of tensors, then applies `torch.cat` along the specified dimension. A single surviving tensor is returned unchanged.

### Parameters
#### Path Parameters
- **tensors** (list of tensors or None) - A list of tensors to concatenate.
- **dim** (int, optional) - The dimension along which to concatenate. Defaults to 0.

### Returns
- The concatenated tensor. Returns `None` when no active tensor remains.
```

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

### save_load decorator

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Class decorator that records constructor arguments and adds save, load, and init_and_load methods to a torch.nn.Module subclass. It simplifies checkpointing by automatically handling the serialization and deserialization of model states and constructor arguments.

```APIDOC
## `save_load(...)` — checkpoint decorator for `nn.Module` subclasses

Class decorator that records constructor arguments and adds `save`, `load`, and `init_and_load` methods to a `torch.nn.Module` subclass.

```python
from pathlib import Path
import torch
from torch import nn, Tensor
from torch_einops_kit.save_load import save_load

@save_load()
class SimpleNet(nn.Module):
    def __init__(self, dim: int, hidden_dim: int) -> None:
        super().__init__()
        self.dim = dim
        self.hidden_dim = hidden_dim
        self.net = nn.Linear(dim, hidden_dim)

    def forward(self, x: Tensor) -> Tensor:
        return self.net(x)

path = Path("model.pt")
model = SimpleNet(10, 20)
model.save(str(path))

# Reconstruct from checkpoint without knowing constructor args
restored = SimpleNet.init_and_load(str(path))
assert restored.dim == 10
assert restored.hidden_dim == 20

# Custom method names and version tagging
@save_load(
    save_method_name="store",
    load_method_name="restore",
    config_instance_var_name="stored_config",
    init_and_load_classmethod_name="create_and_restore",
    version="1.0.0",
)
class VersionedNet(nn.Module):
    def __init__(self, width: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.randn(width))

net = VersionedNet(13)
net.store("versioned_model.pt")
restored_net = VersionedNet.create_and_restore("versioned_model.pt")
assert restored_net.weight.shape == (13,)
```
```

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

### Prefix and Suffix Slicing

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Use `slice_left_at_dim` to retain the first `length` elements along a dimension, and `slice_right_at_dim` to retain the last `length` elements. These are useful for trimming tensors to a specific size, such as when aligning precomputed frequencies to a query length.

```python
import torch
from torch_einops_kit import slice_left_at_dim, slice_right_at_dim

t = torch.randn(3, 4, 5)

prefix = slice_left_at_dim(t, 2, dim=1)
assert prefix.shape == (3, 2, 5)

suffix = slice_right_at_dim(t, 2, dim=1)
assert suffix.shape == (3, 2, 5)

# Typical use: trim precomputed positional frequencies to query length
freqs = torch.randn(1, 512, 64)  # (1, max_seq, head_dim)
q_len = 128
trimmed_freqs = slice_right_at_dim(freqs, q_len, dim=-2)
assert trimmed_freqs.shape == (1, 128, 64)
```

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

### Pack and Unpack Tensors with Inverse Function

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Packs a tensor or list of tensors using an einops pattern and returns a paired inverse function for shape restoration. The inverse function can also accept a modified packed tensor.

```python
import torch
from torch_einops_kit.einops import pack_with_inverse

# Pack a single tensor: flatten spatial dims
t = torch.randn(3, 12, 2, 2)
packed, inverse = pack_with_inverse(t, "b * d")
assert packed.shape == (3, 24, 2)

# Recover original shape
restored = inverse(packed)
assert restored.shape == (3, 12, 2, 2)

# Pack a list of tensors and unpack with an overriding pattern
t = torch.randn(3, 12, 2)
u = torch.randn(3, 4, 2)
packed, inverse = pack_with_inverse([t, u], "b * d")
assert packed.shape == (3, 28, 2)

# Apply an operation on the packed tensor, then unpack
reduced = packed.sum(dim=-1)        # (3, 28)
t_out, u_out = inverse(reduced, "b *")
assert t_out.shape == (3, 12)
assert u_out.shape == (3, 4)
```

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

### pad_sequence

Source: https://context7.com/hunterhogan/torch_einops_kit/llms.txt

Pads a list of tensors to a shared length along a specified dimension and optionally stacks them. It can also return the original lengths and padding widths.

```APIDOC
## pad_sequence(tensors, dim, return_lens, return_stacked, left, pad_lens)

### Description
Pads a sequence of tensors to a shared length along a specified dimension. Options to return original lengths, stack tensors, or return padding widths are available.

### Parameters
#### Path Parameters
- **tensors** (list[torch.Tensor]) - A list of tensors to pad.
- **dim** (int) - The dimension along which to pad.
- **return_lens** (bool, optional) - If True, returns the original lengths of the tensors.
- **return_stacked** (bool, optional) - If True, returns a single stacked tensor. Defaults to True.
- **left** (bool, optional) - If True, performs left padding. Defaults to False.
- **pad_lens** (bool, optional) - If True, returns the padding widths for each sequence. Defaults to False.

### Request Example
```python
import torch
from torch_einops_kit import pad_sequence

x = torch.randn(2, 4, 5)
y = torch.randn(2, 3, 5)
z = torch.randn(2, 1, 5)

# Stack and recover original lengths
packed, lens = pad_sequence([x, y, z], dim=1, return_lens=True)
# packed.shape == (3, 2, 4, 5)
# lens.tolist() == [4, 3, 1]

# Left-pad and retrieve padding widths
prompt_ids = [torch.randint(0, 1000, (t,)) for t in [6, 4, 8]]
Padded_ids, start_pos = pad_sequence(prompt_ids, return_lens=True, left=True, pad_lens=True)
# padded_ids.shape == (3, 8)
# start_pos.tolist() == [2, 4, 0]

# Return a list of individually padded tensors
Padded_list = pad_sequence([x, y, z], dim=1, return_stacked=False)
# isinstance(padded_list, list) == True
# all(t.shape[1] == 4 for t in padded_list) == True
```
```

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

### pad_at_dim

Source: https://github.com/hunterhogan/torch_einops_kit/blob/main/README.md

Pads or trims a tensor along a specified dimension. Positive values add elements, while negative values trim elements.

```APIDOC
## pad_at_dim

### Description
Pads or trims `t` along one dimension. Positive values add elements. Negative values trim elements.

### Parameters
- **t** (Tensor) - The input tensor.
- **pad** (tuple[int, int]) - A tuple of two integers representing the padding for the left and right sides of the dimension.
- **dim** (int, optional) - The dimension along which to pad. Defaults to -1.
- **value** (float, optional) - The padding value. Defaults to 0.0.
```
Tokens:
Raw