Torch Einops Utils

https://github.com/lucidrains/torch-einops-utils
Admin
Torch Einops Utils is a Python utility library providing helper functions for PyTorch tensor...
Tokens:8,908
Snippets:44
Trust Score:10
License:MIT
Update:7 months ago
Show doc for...
Tokens:
Raw
### Install torch-einops-utils

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Install the library via pip.

```bash
pip install torch-einops-utils
```

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

### Pad Sequence with Return Lengths

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Use `pad_sequence` with `return_lens=True` and `pad_lens=True` to get padding lengths instead of sequence lengths. This is useful for understanding the padding applied to sequences.

```python
packed, pad_lens = pad_sequence([x, y, z], dim=1, return_lens=True, pad_lens=True)
print(pad_lens)  # tensor([0, 1, 3])
```

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

### Get Module Device

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Retrieves the device of a PyTorch module. Returns None if the module has no parameters or is not on a device.

```python
from torch_einops_utils.device import module_device

# Module with parameters
model = nn.Linear(3, 3)
device = module_device(model)
print(device)  # device(type='cpu')

# Module without parameters (e.g., Identity)
model = nn.Identity()
device = module_device(model)
print(device)  # None
```

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

### Get PyTorch Module Device

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

The `module_device` function efficiently retrieves the device of a PyTorch module by inspecting its parameters or buffers. This is useful for ensuring tensors are on the correct device before operations.

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

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

### Pad Uneven Images

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Demonstrates padding uneven images along the height dimension and then concatenating them. `pad_sequence` can handle varying dimensions, and `pad_sequence_and_cat` simplifies stacking.

```python
images = [
    torch.randn(3, 16, 17),  # C, H, W
    torch.randn(3, 15, 18),
    torch.randn(3, 17, 16)
]

# Pad height dimension, then concatenate
padded_height = pad_sequence(images, dim=-2, return_stacked=False)
stacked = pad_sequence_and_cat(padded_height, dim_cat=0)
print(stacked.shape)  # torch.Size([9, 17, 18])
```

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

### pad_sequence and pad_sequence_and_cat

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Utilities for padding sequences and images, with options to return lengths or unstacked tensors.

```APIDOC
## pad_sequence and pad_sequence_and_cat

### Description
Utilities for padding sequences and images, with options to return lengths or unstacked tensors.

### Usage Examples

```python
import torch

# Assuming pad_sequence and pad_sequence_and_cat are imported from torch_einops_utils

# Return padding lengths instead of sequence lengths
# packed, pad_lens = pad_sequence([x, y, z], dim=1, return_lens=True, pad_lens=True)
# print(pad_lens)  # tensor([0, 1, 3])

# Without stacking (returns list)
# padded = pad_sequence([x, y, z], dim=1, return_stacked=False)
# print(len(padded))  # 3
# print([t.shape for t in padded])  # All torch.Size([2, 4, 5])

# Padding uneven images
images = [
    torch.randn(3, 16, 17),  # C, H, W
    torch.randn(3, 15, 18),
    torch.randn(3, 17, 16)
]

# Pad height dimension, then concatenate
# padded_height = pad_sequence(images, dim=-2, return_stacked=False)
# stacked = pad_sequence_and_cat(padded_height, dim_cat=0)
# print(stacked.shape)  # torch.Size([9, 17, 18])
```
```

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

### module_device

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Returns the device of a PyTorch module by checking its first parameter or buffer.

```APIDOC
## module_device

### Description
Returns the device of a PyTorch module by checking its first parameter or buffer.

### Method
`module_device(module)`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **module** (nn.Module) - Required - The PyTorch module to check.

### Request Example
```python
import torch
from torch import nn
from torch_einops_utils.device import module_device

# Example usage:
# model = nn.Linear(10, 5).to('cuda')
# device = module_device(model)
# print(device)  # cuda:0
```

### Response
#### Success Response (200)
- **device** (torch.device) - The device of the module.

#### Response Example
```json
{
  "device": "cuda:0"
}
```
```

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

### Pack Tensors with Einops and Inverse

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Use `pack_with_inverse` to pack tensors using einops patterns and obtain an inverse function for unpacking. This simplifies operations requiring reversible reshaping of tensors.

```python
import torch
from torch_einops_utils import pack_with_inverse

# Single tensor packing
t = torch.randn(3, 12, 2, 2)
packed, inverse = pack_with_inverse(t, 'b * d')
print(packed.shape)  # torch.Size([3, 24, 2]) - middle dims merged

# Unpack back to original shape
unpacked = inverse(packed)
print(unpacked.shape)  # torch.Size([3, 12, 2, 2])

# Pack multiple tensors
t = torch.randn(3, 12, 2, 2)
u = torch.randn(3, 4, 2)
packed, inverse = pack_with_inverse([t, u], 'b * d')
print(packed.shape)  # torch.Size([3, 28, 2])

# Unpack with different pattern
packed_sum = packed.sum(dim=-1)
t_out, u_out = inverse(packed_sum, 'b *')
print(t_out.shape)  # torch.Size([3, 12, 2])
print(u_out.shape)  # torch.Size([3, 4])
```

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

### pack_with_inverse

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Packs tensors using einops and returns an inverse function to unpack them. Simplifies reshape operations that need to be reversed.

```APIDOC
## pack_with_inverse

### Description
Packs tensors using einops and returns an inverse function to unpack them. Simplifies reshape operations that need to be reversed.

### Method
`pack_with_inverse(tensors, pattern)`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **tensors** (torch.Tensor or list of torch.Tensor) - Required - The tensor(s) to pack.
- **pattern** (str) - Required - The einops pattern for packing.

### Request Example
```python
import torch
from torch_einops_utils import pack_with_inverse

# Single tensor packing
t = torch.randn(3, 12, 2, 2)
packed, inverse = pack_with_inverse(t, 'b * d')
print(packed.shape)  # torch.Size([3, 24, 2]) - middle dims merged

# Unpack back to original shape
unpacked = inverse(packed)
print(unpacked.shape)  # torch.Size([3, 12, 2, 2])

# Pack multiple tensors
t = torch.randn(3, 12, 2, 2)
u = torch.randn(3, 4, 2)
packed, inverse = pack_with_inverse([t, u], 'b * d')
print(packed.shape)  # torch.Size([3, 28, 2])

# Unpack with different pattern
packed_sum = packed.sum(dim=-1)
t_out, u_out = inverse(packed_sum, 'b *')
print(t_out.shape)  # torch.Size([3, 12, 2])
print(u_out.shape)  # torch.Size([3, 4])
```

### Response
#### Success Response (200)
- **packed** (torch.Tensor) - The packed tensor.
- **inverse** (function) - A function to unpack the tensor(s).

#### Response Example
```json
{
  "packed": "<torch.Tensor shape=[3, 24, 2]>",
  "inverse": "<function pack_with_inverse.<locals>.inverse at 0x...>"
}
```
```

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

### Handle None inputs with maybe

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Wraps a function to return None if the input is None, preventing errors during chained operations.

```python
import torch
from torch_einops_utils import maybe

# Create a function that adds 1 to a tensor
add_one = lambda t: t + 1

# Wrap it with maybe to handle None safely
safe_add_one = maybe(add_one)

# Works normally with tensors
tensor = torch.tensor([1, 2, 3])
result = safe_add_one(tensor)
print(result)  # tensor([2, 3, 4])

# Returns None when input is None (no error)
result = safe_add_one(None)
print(result)  # None

# maybe(None) returns identity function
identity_fn = maybe(None)
result = identity_fn(tensor)
print(result)  # tensor([1, 2, 3])
```

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

### Pad Sequence Without Stacking

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

When `return_stacked` is `False`, `pad_sequence` returns a list of padded tensors. This is useful when you need to process padded sequences individually.

```python
padded = pad_sequence([x, y, z], dim=1, return_stacked=False)
print(len(padded))  # 3
print([t.shape for t in padded])  # All torch.Size([2, 4, 5])
```

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

### maybe - Null-Safe Function Wrapper

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Wraps a function to handle None inputs gracefully. If the input is None, returns None instead of calling the function. Useful for chaining operations on potentially missing tensors.

```APIDOC
## maybe - Null-Safe Function Wrapper

### Description
Wraps a function to handle None inputs gracefully. If the input is None, returns None instead of calling the function. Useful for chaining operations on potentially missing tensors.

### Usage
```python
from torch_einops_utils import maybe

# Example function
add_one = lambda t: t + 1

# Wrap with maybe
safe_add_one = maybe(add_one)

# Usage with tensor
# result = safe_add_one(tensor)

# Usage with None
# result = safe_add_one(None) # returns None

# maybe(None) returns identity function
# identity_fn = maybe(None)
# result = identity_fn(tensor)
```

### Parameters
- **fn** (callable) - The function to wrap.

### Returns
- (callable) - A new function that handles None inputs.
```

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

### Move Inputs to Module Device Decorator

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Decorator for nn.Module methods that automatically moves tensor inputs to the same device as the module's parameters. Ensures inputs and model are on the same device.

```python
import torch
from torch import nn
from torch_einops_utils.device import move_inputs_to_module_device

class MyModule(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.linear = nn.Linear(dim, dim)

    @move_inputs_to_module_device
    def forward(self, x, mask=None):
        out = self.linear(x)
        if mask is not None:
            out = out * mask
        return out

model = MyModule(64)
x = torch.randn(2, 64)
output = model(x)
```

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

### Move Inputs to Specific Device Decorator

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Decorator to automatically move all tensor inputs of a function to a specified device before execution. Useful for ensuring consistent device placement.

```python
import torch
from torch import nn
from torch_einops_utils.device import move_inputs_to_device

@move_inputs_to_device(torch.device('cpu'))
def process(x, y):
    return x + y

x = torch.randn(3)
y = torch.randn(3)
result = process(x, y)
```

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

### Pad and Concatenate Variable-Length Sequences

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Pads a list of tensors with varying lengths to a uniform size, optionally stacking them. Supports both right and left padding, with left padding being useful for causal models.

```python
import torch
from torch_einops_utils import pad_sequence, pad_sequence_and_cat

# Variable length sequences
x = torch.randn(2, 4, 5)  # batch=2, seq=4, features=5
y = torch.randn(2, 3, 5)  # batch=2, seq=3, features=5
z = torch.randn(2, 1, 5)  # batch=2, seq=1, features=5

# Pad and stack along sequence dimension
packed, lens = pad_sequence([x, y, z], dim=1, return_lens=True)
print(packed.shape)  # torch.Size([3, 2, 4, 5]) - stacked with batch dim first
print(lens)  # tensor([4, 3, 1])

# Left padding (useful for causal/autoregressive models)
packed = pad_sequence([x, y, z], dim=1, left=True)
print(packed.shape)  # torch.Size([3, 2, 4, 5])
```

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

### Pad tensor dimensions

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Adds singleton dimensions to the left or right of a tensor to facilitate broadcasting.

```python
import torch
from torch_einops_utils import (
    pad_ndim, pad_left_ndim, pad_right_ndim,
    pad_left_ndim_to, pad_right_ndim_to, align_dims_left
)

t = torch.randn(3)

# Add 1 dim left and 2 dims right
result = pad_ndim(t, (1, 2))
print(result.shape)  # torch.Size([1, 3, 1, 1])

# Add dimensions only to right
result = pad_right_ndim(t, 2)
print(result.shape)  # torch.Size([3, 1, 1])

# Add dimensions only to left
result = pad_left_ndim(t, 2)
print(result.shape)  # torch.Size([1, 1, 3])
```

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

### save_load Decorator

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

A decorator that adds `save`, `load`, and `init_and_load` methods to PyTorch modules. Automatically preserves constructor arguments for full model reconstruction.

```APIDOC
## save_load Decorator

### Description
A decorator that adds `save`, `load`, and `init_and_load` methods to PyTorch modules. Automatically preserves constructor arguments for full model reconstruction.

### Usage
Apply the `@save_load()` decorator to your `nn.Module` class.

### Methods Added
- **`save(path)`**: Saves the model state and constructor arguments to the specified path.
- **`load(path)`**: Loads the model state from the specified path into the current instance.
- **`init_and_load(path)`**: Creates a new instance of the model using saved constructor arguments and loads the state from the specified path.

### Request Example
```python
import torch
from torch import nn
from torch_einops_utils.save_load import save_load

@save_load(version='1.0.0')
class MyModel(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.1):
        super().__init__()
        self.dim = dim
        self.hidden_dim = hidden_dim
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim)
        )

    def forward(self, x):
        return self.net(x)

# Create and save model
model = MyModel(dim=64, hidden_dim=256, dropout=0.2)
model.save('model.pt')

# Load into existing model
model2 = MyModel(dim=64, hidden_dim=256, dropout=0.2)
model2.load('model.pt')

# Initialize AND load from checkpoint (reconstructs with saved config)
model3 = MyModel.init_and_load('model.pt')
print(model3.dim)  # 64
print(model3.hidden_dim)  # 256

# Works with nested decorated modules
@save_load()
class Encoder(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.proj = nn.Linear(dim, dim)

@save_load()
class FullModel(nn.Module):
    def __init__(self, encoder, output_dim):
        super().__init__()
        self.encoder = encoder
        self.output = nn.Linear(encoder.dim, output_dim)

# Nested models are automatically serialized
encoder = Encoder(dim=32)
model = FullModel(encoder=encoder, output_dim=10)
model.save('full_model.pt')

# Reconstruct entire hierarchy
loaded = FullModel.init_and_load('full_model.pt')
print(loaded.encoder.dim)  # 32
```

### Response
#### Success Response (200)
- **Model Instance**: The decorated model class gains `save`, `load`, and `init_and_load` methods.

#### Response Example
(No direct response body, but methods are added to the decorated class.)
```

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

### Combine Boolean Masks with AND/OR

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Combine multiple boolean masks using logical AND or OR operations. Safely handles None values in the input list, returning None if all masks are None.

```python
import torch
from torch_einops_utils import and_masks, or_masks

mask1 = torch.tensor([True, True, False, False])
mask2 = torch.tensor([True, False, True, False])

# AND combination
result = and_masks([mask1, mask2])
print(result)  # tensor([True, False, False, False])

# OR combination
result = or_masks([mask1, mask2])
print(result)  # tensor([True, True, True, False])

# None values are ignored
result = and_masks([mask1, None, mask2])
print(result)  # tensor([True, False, False, False])

# Returns None if all masks are None
result = and_masks([None, None])
print(result)  # None
```

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

### Slice tensors at specific dimensions

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Provides utility functions for slicing tensors at arbitrary dimensions using slice objects or length-based offsets.

```python
import torch
from torch_einops_utils import slice_at_dim, slice_left_at_dim, slice_right_at_dim

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

# Slice at last dimension (default dim=-1)
result = slice_at_dim(t, slice(1, 3))
print(result.shape)  # torch.Size([3, 4, 2])

# Slice at specific dimension
result = slice_at_dim(t, slice(None, 2), dim=1)
print(result.shape)  # torch.Size([3, 2, 5])

# Slice from start (left) at dimension
result = slice_left_at_dim(t, length=2, dim=1)
print(result.shape)  # torch.Size([3, 2, 5])
# Equivalent to t[:, :2, :]

# Slice from end (right) at dimension
result = slice_right_at_dim(t, length=2, dim=1)
print(result.shape)  # torch.Size([3, 2, 5])
# Equivalent to t[:, -2:, :]
```

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

### Model Serialization with save_load Decorator

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

The `@save_load` decorator adds `save`, `load`, and `init_and_load` methods to PyTorch modules, automatically handling constructor arguments for complete model reconstruction. It supports nested decorated modules.

```python
import torch
from torch import nn
from torch_einops_utils.save_load import save_load

@save_load(version='1.0.0')
class MyModel(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.1):
        super().__init__()
        self.dim = dim
        self.hidden_dim = hidden_dim
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim)
        )

    def forward(self, x):
        return self.net(x)

# Create and save model
model = MyModel(dim=64, hidden_dim=256, dropout=0.2)
model.save('model.pt')

# Load into existing model
model2 = MyModel(dim=64, hidden_dim=256, dropout=0.2)
model2.load('model.pt')

# Initialize AND load from checkpoint (reconstructs with saved config)
model3 = MyModel.init_and_load('model.pt')
print(model3.dim)  # 64
print(model3.hidden_dim)  # 256

# Works with nested decorated modules
@save_load()
class Encoder(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.proj = nn.Linear(dim, dim)

@save_load()
class FullModel(nn.Module):
    def __init__(self, encoder, output_dim):
        super().__init__()
        self.encoder = encoder
        self.output = nn.Linear(encoder.dim, output_dim)

# Nested models are automatically serialized
encoder = Encoder(dim=32)
model = FullModel(encoder=encoder, output_dim=10)
model.save('full_model.pt')

# Reconstruct entire hierarchy
loaded = FullModel.init_and_load('full_model.pt')
print(loaded.encoder.dim)  # 32
```

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

### masked_mean - Compute Mean with Masking Support

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Computes the mean of a tensor with optional masking. Supports broadcasting masks to match tensor dimensions and handles edge cases where mask is all False.

```APIDOC
## masked_mean - Compute Mean with Masking Support

### Description
Computes the mean of a tensor with optional masking. Supports broadcasting masks to match tensor dimensions and handles edge cases where mask is all False.

### Usage
```python
import torch
from torch_einops_utils import masked_mean

# Example tensor and mask
t = torch.tensor([1., 2., 3., 4.])
mask = torch.tensor([True, False, True, False])

# Compute mean without mask
# result = masked_mean(t)

# Compute mean with mask
# result = masked_mean(t, mask=mask)

# Compute mean along a dimension with mask
# t_2d = torch.tensor([[1., 2.], [3., 4.]])
# mask_2d = torch.tensor([[True, False], [True, True]])
# result_dim0 = masked_mean(t_2d, mask=mask_2d, dim=0)
# result_dim1 = masked_mean(t_2d, mask=mask_2d, dim=1)

# Mask with fewer dimensions is auto-expanded
# t_3d = torch.randn(2, 3, 4)
# mask_1d = torch.tensor([True, False])
# result_3d = masked_mean(t_3d, mask=mask_1d, dim=(1, 2))
```

### Parameters
- **tensor** (torch.Tensor) - The input tensor.
- **mask** (torch.Tensor, optional) - A boolean mask. Defaults to None.
- **dim** (int or tuple of ints, optional) - The dimension(s) along which to compute the mean. Defaults to None (computes mean over all dimensions).

### Returns
- (torch.Tensor) - The computed mean.
```

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

### Pad Tensor at Specific Dimension

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Pad tensors at any dimension with specified values. Supports left, right, or both-sided padding, and padding to a target length. Custom padding values can be used.

```python
import torch
from torch_einops_utils import (
    pad_at_dim, pad_left_at_dim, pad_right_at_dim,
    pad_left_at_dim_to, pad_right_at_dim_to
)

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

# Pad with (left_pad, right_pad) at dimension 1
result = pad_at_dim(t, (0, 1), dim=1)
print(result.shape)  # torch.Size([3, 7, 1])

# Pad only on the right at dimension 1
result = pad_right_at_dim(t, 2, dim=1)
print(result.shape)  # torch.Size([3, 8, 1])

# Pad only on the left at dimension 1
result = pad_left_at_dim(t, 2, dim=1)
print(result.shape)  # torch.Size([3, 8, 1])

# Pad to reach a target length
result = pad_right_at_dim_to(t, length=10, dim=1)
print(result.shape)  # torch.Size([3, 10, 1])

result = pad_left_at_dim_to(t, length=10, dim=1)
print(result.shape)  # torch.Size([3, 10, 1])

# Custom padding value
result = pad_at_dim(t, (1, 1), dim=1, value=-1.0)
print(result[:, 0, :])  # All -1.0
```

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

### Safe Stack and Concatenate Tensors

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Stack or concatenate tensors while safely handling None values. Returns None for empty lists, or the original tensor if only one element is present. None values in the list are ignored.

```python
import torch
from torch_einops_utils import safe_stack, safe_cat

t1 = torch.randn(2, 3)
t2 = torch.randn(2, 3)

# Safe stack - handles None values
result = safe_stack([t1, t2])
print(result.shape)  # torch.Size([2, 2, 3])

result = safe_stack([t1, None, t2])
print(result.shape)  # torch.Size([2, 2, 3]) - None ignored

result = safe_stack([t1])
print(result.shape)  # torch.Size([1, 2, 3])

result = safe_stack([])
print(result)  # None

result = safe_stack([None, None])
print(result)  # None

# Safe cat - similar behavior
result = safe_cat([t1, t2])
print(result.shape)  # torch.Size([4, 3])

result = safe_cat([t1, None])
print(result == t1)  # All True - returns single tensor as-is

result = safe_cat([])
print(result)  # None
```

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

### Pad Tensor to Specific Number of Dimensions

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Use pad_right_ndim_to and pad_left_ndim_to to add dimensions to a tensor by padding on the right or left, respectively. Useful for broadcasting or aligning tensors.

```python
import torch
from torch_einops_utils import pad_right_ndim_to, pad_left_ndim_to

t = torch.randn(3)
result = pad_right_ndim_to(t, 3)
print(result.shape)  # torch.Size([3, 1, 1])

result = pad_left_ndim_to(t, 3)
print(result.shape)  # torch.Size([1, 1, 3])

t = torch.randn(3)
u = torch.randn(3, 5, 2)
v = torch.randn(3, 5)

t, u, v = align_dims_left((t, u, v))
print(t.shape)  # torch.Size([3, 1, 1])
print(u.shape)  # torch.Size([3, 5, 2])
print(v.shape)  # torch.Size([3, 5, 1])
```

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

### Convert Sequence Lengths to Boolean Mask

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Generates a boolean mask from a tensor of sequence lengths. Useful for masking variable-length sequences in batches. Can specify a maximum length and works with batched lengths.

```python
import torch
from torch_einops_utils import lens_to_mask

# Sequence lengths
lens = torch.tensor([3, 2, 4])

# Create mask where True indicates valid positions
mask = lens_to_mask(lens)
print(mask)
# tensor([[ True,  True,  True, False],
#         [ True,  True, False, False],
#         [ True,  True,  True,  True]])

# Specify maximum length explicitly
mask = lens_to_mask(lens, max_len=6)
print(mask.shape)  # torch.Size([3, 6])

# Works with batched lengths
lens = torch.tensor([[2, 3], [1, 4]])
mask = lens_to_mask(lens)
print(mask.shape)  # torch.Size([2, 2, 4])
```

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

### pad_ndim - Add Dimensions to Tensors

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Adds singleton dimensions to the left and/or right of a tensor's shape. Essential for broadcasting operations.

```APIDOC
## pad_ndim - Add Dimensions to Tensors

### Description
Adds singleton dimensions to the left and/or right of a tensor's shape. Essential for broadcasting operations.

### Usage
```python
import torch
from torch_einops_utils import (
    pad_ndim, pad_left_ndim, pad_right_ndim,
    pad_left_ndim_to, pad_right_ndim_to, align_dims_left
)

t = torch.randn(3)

# Add 1 dim left and 2 dims right
# result = pad_ndim(t, (1, 2))

# Add dimensions only to right
# result = pad_right_ndim(t, 2)

# Add dimensions only to left
# result = pad_left_ndim(t, 2)
```

### Parameters
- **tensor** (torch.Tensor) - The input tensor.
- **pad_left** (int, optional) - Number of dimensions to add to the left. Defaults to 0.
- **pad_right** (int, optional) - Number of dimensions to add to the right. Defaults to 0.

### Returns
- (torch.Tensor) - The tensor with added dimensions.
```

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

### tree_flatten_with_inverse

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Flattens a nested Python structure (PyTree) and returns an inverse function to reconstruct the original structure.

```APIDOC
## tree_flatten_with_inverse

### Description
Flattens a nested Python structure (PyTree) and returns an inverse function to reconstruct the original structure.

### Method
`tree_flatten_with_inverse(tree)`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **tree** (PyTree) - Required - The nested Python structure to flatten.

### Request Example
```python
import torch
from torch_einops_utils import tree_flatten_with_inverse

# Nested structure
tree = (1, (2, 3), 4)

# Flatten with inverse
flattened, inverse = tree_flatten_with_inverse(tree)
print(flattened)  # [1, 2, 3, 4]

# Modify and reconstruct
modified = [flattened[0] + 10] + flattened[1:]
reconstructed = inverse(modified)
print(reconstructed)  # (11, (2, 3), 4)

# Works with complex nested structures
nested = {'a': [1, 2], 'b': (3, {'c': 4})}
flattened, inverse = tree_flatten_with_inverse(nested)
reconstructed = inverse(flattened)
print(reconstructed == nested)  # True
```

### Response
#### Success Response (200)
- **flattened** (list) - A flat list of all elements in the PyTree.
- **inverse** (function) - A function that takes a list of elements and reconstructs the original PyTree structure.

#### Response Example
```json
{
  "flattened": [1, 2, 3, 4],
  "inverse": "<function tree_flatten_with_inverse.<locals>.inverse at 0x...>"
}
```
```

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

### slice_at_dim - Slice Tensor at Any Dimension

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Slices a tensor at a specified dimension using Python slice objects. Provides cleaner syntax than manual indexing for variable dimensions.

```APIDOC
## slice_at_dim - Slice Tensor at Any Dimension

### Description
Slices a tensor at a specified dimension using Python slice objects. Provides cleaner syntax than manual indexing for variable dimensions.

### Usage
```python
import torch
from torch_einops_utils import slice_at_dim, slice_left_at_dim, slice_right_at_dim

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

# Slice at last dimension (default dim=-1)
# result = slice_at_dim(t, slice(1, 3))

# Slice at specific dimension
# result = slice_at_dim(t, slice(None, 2), dim=1)

# Slice from start (left) at dimension
# result = slice_left_at_dim(t, length=2, dim=1)

# Slice from end (right) at dimension
# result = slice_right_at_dim(t, length=2, dim=1)
```

### Parameters
- **tensor** (torch.Tensor) - The input tensor.
- **slc** (slice) - The slice object to apply.
- **dim** (int, optional) - The dimension to slice. Defaults to -1 (last dimension).

### Returns
- (torch.Tensor) - The sliced tensor.
```

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

### Flatten PyTree with Inverse Reconstruction

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Use `tree_flatten_with_inverse` to flatten a nested Python structure (PyTree) into a list and obtain a function to reconstruct the original structure. This is useful for modifying nested data and then restoring its shape.

```python
import torch
from torch_einops_utils import tree_flatten_with_inverse

# Nested structure
tree = (1, (2, 3), 4)

# Flatten with inverse
flattened, inverse = tree_flatten_with_inverse(tree)
print(flattened)  # [1, 2, 3, 4]

# Modify and reconstruct
modified = [flattened[0] + 10] + flattened[1:]
reconstructed = inverse(modified)
print(reconstructed)  # (11, (2, 3), 4)

# Works with complex nested structures
nested = {'a': [1, 2], 'b': (3, {'c': 4})}
flattened, inverse = tree_flatten_with_inverse(nested)
reconstructed = inverse(flattened)
print(reconstructed == nested)  # True
```

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

### Apply Function to Tensors in PyTree

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

The `tree_map_tensor` function applies a given function exclusively to tensors within a PyTree, leaving non-tensor elements untouched. This is ideal for element-wise operations on tensor components of complex data structures.

```python
import torch
from torch_einops_utils import tree_map_tensor

# Mixed structure with tensors and non-tensors
tree = (1, torch.tensor([2, 3]), 'hello', torch.tensor([4.0]))

# Apply function only to tensors
result = tree_map_tensor(lambda t: t * 2, tree)
print(result[0])  # 1 (unchanged - not a tensor)
print(result[1])  # tensor([4, 6])
print(result[2])  # 'hello' (unchanged)
print(result[3])  # tensor([8.0])
```

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

### Compute masked mean

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Calculates the mean of a tensor while respecting a boolean mask, supporting broadcasting and specific dimensions.

```python
import torch
from torch_einops_utils import masked_mean

# Simple mean without mask
t = torch.tensor([1., 2., 3., 4.])
result = masked_mean(t)
print(result)  # tensor(2.5)

# Mean with a boolean mask
mask = torch.tensor([True, False, True, False])
result = masked_mean(t, mask=mask)
print(result)  # tensor(2.0) - only values at True positions

# Mean along specific dimension with mask
t = torch.tensor([[1., 2.], [3., 4.]])
mask = torch.tensor([[True, False], [True, True]])
result = masked_mean(t, mask=mask, dim=0)
print(result)  # tensor([2.0, 4.0])

result = masked_mean(t, mask=mask, dim=1)
print(result)  # tensor([1.0, 3.5])

# Mask with fewer dimensions is auto-expanded
t = torch.randn(2, 3, 4)
mask = torch.tensor([True, False])  # Shape (2,)
result = masked_mean(t, mask=mask, dim=(1, 2))
print(result.shape)  # torch.Size([2])
```

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

### Modify tensor shapes

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Generates a new shape tuple by replacing specific dimension sizes.

```python
import torch
from torch_einops_utils import shape_with_replace

t = torch.randn(3, 4, 5)
print(t.shape)  # torch.Size([3, 4, 5])

# Replace dimension 1 with value 2
new_shape = shape_with_replace(t, {1: 2})
print(new_shape)  # torch.Size([3, 2, 5])

# Multiple replacements
new_shape = shape_with_replace(t, {0: 1, 2: 10})
print(new_shape)  # torch.Size([1, 4, 10])
```

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

### tree_map_tensor

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Applies a function only to tensors within a nested Python structure, leaving other values unchanged.

```APIDOC
## tree_map_tensor

### Description
Applies a function only to tensors within a nested Python structure, leaving other values unchanged.

### Method
`tree_map_tensor(func, tree)`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **func** (function) - Required - The function to apply to each tensor.
- **tree** (PyTree) - Required - The nested Python structure to process.

### Request Example
```python
import torch
from torch_einops_utils import tree_map_tensor

# Mixed structure with tensors and non-tensors
tree = (1, torch.tensor([2, 3]), 'hello', torch.tensor([4.0]))

# Apply function only to tensors
result = tree_map_tensor(lambda t: t * 2, tree)
print(result[0])  # 1 (unchanged - not a tensor)
print(result[1])  # tensor([4, 6])
print(result[2])  # 'hello' (unchanged)
print(result[3])  # tensor([8.0])
```

### Response
#### Success Response (200)
- **result** (PyTree) - The PyTree with the function applied to its tensor elements.

#### Response Example
```python
# Example output structure (actual tensor values will differ)
(1, tensor([4, 6]), 'hello', tensor([8.0]))
```
```

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

### shape_with_replace - Modify Tensor Shape

Source: https://context7.com/lucidrains/torch-einops-utils/llms.txt

Returns a modified shape by replacing values at specified indices. Useful for creating new shapes based on existing tensors.

```APIDOC
## shape_with_replace - Modify Tensor Shape

### Description
Returns a modified shape by replacing values at specified indices. Useful for creating new shapes based on existing tensors.

### Usage
```python
import torch
from torch_einops_utils import shape_with_replace

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

# Replace dimension 1 with value 2
# new_shape = shape_with_replace(t, {1: 2})

# Multiple replacements
# new_shape = shape_with_replace(t, {0: 1, 2: 10})
```

### Parameters
- **tensor** (torch.Tensor) - The input tensor.
- **replacements** (dict) - A dictionary where keys are dimension indices and values are the new sizes for those dimensions.

### Returns
- (torch.Size) - The new shape.
```
Tokens:
Raw