torch-einops-kit 0.1.2__tar.gz → 0.1.3__tar.gz

{torch_einops_kit-0.1.2 → torch_einops_kit-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-einops-kit
-Version: 0.1.2
+Version: 0.1.3
 Summary: Typed tensor-shaping, masking, padding, device-routing, and checkpoint utilities for PyTorch and `einops`. A superset of `lucidrains/torch-einops-utils` with similar utilities from other lucidrains packages. Adds strict typing, extensive tests, and comprehensive docstrings.
 Keywords: artificial intelligence,einops,machine learning,pytorch,torch
 Author: Phil Wang, Hunter Hogan
@@ -19,12 +19,12 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Utilities
 Classifier: Typing :: Typed
 Requires-Dist: einops>=0.8.2
-Requires-Dist: torch>=2.10.0
+Requires-Dist: torch>=2.0.0
 Requires-Dist: typing-extensions>=4.0.0
 Maintainer: Hunter Hogan
 Maintainer-email: Hunter Hogan <HunterHogan@pm.me>
 Requires-Python: >=3.10
-Project-URL: Documentation, https://context7.com/hunterhogan/torch_einops_kit
+Project-URL: Context7, https://context7.com/hunterhogan/torch_einops_kit
 Project-URL: Donate, https://www.patreon.com/integrated
 Project-URL: Download, https://pypi.org/project/torch_einops_kit
 Project-URL: Homepage, https://github.com/hunterhogan/torch_einops_kit
@@ -34,7 +34,7 @@ Description-Content-Type: text/markdown
 
 # torch_einops_kit
 
-Typed tensor-shaping, masking, padding, device-routing, and checkpoint utilities for PyTorch and `einops`.
+Typed tensor-shaping, masking, padding, device-routing, lightweight `nn.Module` adapters, and checkpoint utilities for PyTorch and `einops`.
 
 [![pip install torch-einops-kit](https://img.shields.io/badge/pip_install-torch--einops--kit-gray.svg?labelColor=blue)](https://pypi.org/project/torch-einops-kit/)
 [![uv add torch-einops-kit](https://img.shields.io/badge/uv_add-torch--einops--kit-gray.svg?labelColor=blue)](https://pypi.org/project/torch-einops-kit/)
@@ -43,7 +43,7 @@ This repository is a superset of [`lucidrains/torch-einops-utils`](https://githu
 
 `torch_einops_kit` is most useful when combined with other lucidrains repositories. Repositories such as [`dreamer4`](https://github.com/lucidrains/dreamer4), [`metacontroller`](https://github.com/lucidrains/metacontroller), [`mimic-video`](https://github.com/lucidrains/mimic-video), [`pi-zero-pytorch`](https://github.com/lucidrains/pi-zero-pytorch), [`sdft-pytorch`](https://github.com/lucidrains/sdft-pytorch), and [`locoformer`](https://github.com/lucidrains/locoformer) repeatedly need operations such as `align_dims_left`, `shape_with_replace`, `lens_to_mask`, `pad_sequence`, `safe_cat`, and `pack_with_inverse`. This package centralizes those operations in one typed import surface instead of re-implementing the same tensor utility layer in each model repository.
 
-If you already know `torch-einops-utils`, `torch_einops_kit` began as a typed substitute for that package and has since grown into a superset. In addition to everything from `torch-einops-utils`, this repository centralizes small utility functions that appear repeatedly in other lucidrains model repositories but were never collected in one place, such as `l2norm`, `once`, `pack_one`, and `unpack_one`. The function family remains the same kind: small PyTorch and `einops` helpers for shape work, masks, padding, optional tensors, PyTree traversal, device routing, and checkpoint reconstruction. The import path is `torch_einops_kit`, not `torch_einops_utils`. The relationship is conceptual, not literal import-path compatibility.
+If you already know `torch-einops-utils`, `torch_einops_kit` began as a typed substitute for that package and has since grown into a superset. In addition to everything from `torch-einops-utils`, this repository centralizes small utility functions that appear repeatedly in other lucidrains model repositories but were never collected in one place, such as `l2norm`, `once`, `pack_one`, and `unpack_one`. The function family remains the same kind: small PyTorch and `einops` helpers for shape work, masks, padding, optional tensors, PyTree traversal, device routing, lightweight `nn.Module` adaptation, and checkpoint reconstruction. The import path is `torch_einops_kit`, not `torch_einops_utils`. The relationship is conceptual, not literal import-path compatibility.
 
 Use `torch_einops_kit` when you want strict typing, a `py.typed` marker, focused modules, extensive tests, and docstrings written for both humans and AI assistants. Use upstream when you want the most compact possible version of the same idea.
 
@@ -54,7 +54,7 @@ Use `torch_einops_kit` when you want strict typing, a `py.typed` marker, focused
 - Python requirement: `>=3.10`.
 - Runtime dependencies: `torch`, `einops`, and `typing-extensions`.
 - Root package exports: helper functions, slicing helpers, rank-alignment helpers, mask helpers, safe concatenation helpers, padding helpers, normalization helpers, and PyTree / `einops` helpers.
-- Submodules with dedicated imports: `torch_einops_kit.device`, `torch_einops_kit.einops`, `torch_einops_kit.save_load`, and `torch_einops_kit.scaleValues`.
+- Submodules with dedicated imports: `torch_einops_kit.device`, `torch_einops_kit.einops`, `torch_einops_kit.nn`, `torch_einops_kit.save_load`, and `torch_einops_kit.scaleValues`.
 - Typing status: the package ships a `py.typed` marker and the repository uses strict type checking.
 - Best fit: lucidrains-style model repositories that work with variable-length tensors, `einops` patterns, optional intermediate tensors, and nested `torch.nn.Module` graphs.
 
@@ -116,6 +116,16 @@ from torch_einops_kit.device import (
 )
 ```
 
+Import lightweight `nn.Module` adapters from `torch_einops_kit.nn`:
+
+```python
+from torch_einops_kit.nn import (
+    Identity,
+    Lambda,
+    Sequential,
+)
+```
+
 Import checkpoint decorators from `torch_einops_kit.save_load`:
 
 ```python
@@ -126,7 +136,7 @@ from torch_einops_kit.save_load import (
 )
 ```
 
-Import checkpoint decorators from `torch_einops_kit.scaleValues`:
+Import normalization and masked-reduction helpers from `torch_einops_kit.scaleValues`:
 
 ```python
 from torch_einops_kit.scaleValues import (
@@ -296,10 +306,10 @@ When `pad_lens=True` and `return_lens=True`, the second tensor contains padding
 
 ### PyTree helpers
 
-| Name                              | Contract                                                                                                                                                               |
-| --------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `tree_map_tensor(fn, tree)`       | Applies `fn` to every tensor leaf in a PyTree and leaves non-tensor leaves unchanged.                                                                                  |
-| `tree_flatten_with_inverse(tree)` | Returns a flat list of leaves and an inverse function that reconstructs the original PyTree shape from a replacement iterable of leaves.                               |
+| Name                              | Contract                                                                                                                                 |
+| --------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
+| `tree_map_tensor(fn, tree)`       | Applies `fn` to every tensor leaf in a PyTree and leaves non-tensor leaves unchanged.                                                    |
+| `tree_flatten_with_inverse(tree)` | Returns a flat list of leaves and an inverse function that reconstructs the original PyTree shape from a replacement iterable of leaves. |
 
 ## `scaleValues` submodule reference
 
@@ -332,6 +342,16 @@ The `torch_einops_kit.device` submodule contains three utilities for device infe
 | `move_inputs_to_device(device)`    | Decorator that recursively moves every tensor inside positional and keyword arguments to `device` before calling the wrapped function. Non-tensor values pass through unchanged.                                                                                |
 | `move_inputs_to_module_device(fn)` | 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. |
 
+## `nn` submodule reference
+
+The `torch_einops_kit.nn` submodule contains lightweight `torch.nn.Module` adapters and a `Sequential` constructor that ignores `None` values.
+
+| Name                   | Contract                                                                                                                                                                                                       |
+| ---------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `Sequential(*modules)` | Equivalent to `nn.Sequential(*compact(modules))`. Every `None` argument is discarded before construction.                                                                                                      |
+| `Identity()`           | `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. |
+| `Lambda(fn)`           | `torch.nn.Module` wrapper that stores `fn` on `self.fn` and delegates `forward(*args, **kwargs)` to `fn(*args, **kwargs)` without changing the argument structure.                                             |
+
 ## `save_load` submodule reference
 
 The `torch_einops_kit.save_load` submodule contains the checkpoint decorator and the two advanced configuration helpers that support nested decorated modules.
@@ -406,7 +426,7 @@ If you are an AI assistant adapting code from `torch-einops-utils`, use these tr
 This repository is not a repackaged mirror of upstream. This repository makes a different trade-off.
 
 - Upstream is intentionally compact.
-- This fork splits the implementation across focused modules such as `_helpers.py`, `_padding.py`, `device.py`, and `save_load.py` while still re-exporting most tensor helpers from the package root.
+- This fork splits the implementation across focused modules such as `_helpers.py`, `_padding.py`, `device.py`, `nn.py`, and `save_load.py` while still re-exporting most tensor helpers from the package root.
 - This fork adds strict typing, a `py.typed` marker, extensive tests, and detailed docstrings.
 - This fork is best treated as a typed, documented branch of the same utility idea rather than a literal import-path-compatible drop-in replacement.
 
@@ -414,6 +434,7 @@ This repository is not a repackaged mirror of upstream. This repository makes a
 
 - `src/torch_einops_kit/` — package source.
 - `src/torch_einops_kit/device.py` — device inference and input-routing decorators.
+- `src/torch_einops_kit/nn.py` — lightweight `torch.nn.Module` adapters and a `Sequential` constructor that ignores `None` modules.
 - `src/torch_einops_kit/save_load.py` — checkpoint save / load decorator and nested reconstruction helpers.
 - `src/torch_einops_kit/scaleValues.py` — vector normalization, masked mean, and the `RMSNorm` layer.
 - `tests/` — regression tests and usage examples for helpers, masks, padding, device routing, and checkpoint reconstruction.
@@ -435,7 +456,6 @@ pytest
 Run static analysis:
 
 ```bash
-pyright
 ruff check .
 ```
 

{torch_einops_kit-0.1.2 → torch_einops_kit-0.1.3}/README.md

@@ -1,6 +1,6 @@
 # torch_einops_kit
 
-Typed tensor-shaping, masking, padding, device-routing, and checkpoint utilities for PyTorch and `einops`.
+Typed tensor-shaping, masking, padding, device-routing, lightweight `nn.Module` adapters, and checkpoint utilities for PyTorch and `einops`.
 
 [![pip install torch-einops-kit](https://img.shields.io/badge/pip_install-torch--einops--kit-gray.svg?labelColor=blue)](https://pypi.org/project/torch-einops-kit/)
 [![uv add torch-einops-kit](https://img.shields.io/badge/uv_add-torch--einops--kit-gray.svg?labelColor=blue)](https://pypi.org/project/torch-einops-kit/)
@@ -9,7 +9,7 @@ This repository is a superset of [`lucidrains/torch-einops-utils`](https://githu
 
 `torch_einops_kit` is most useful when combined with other lucidrains repositories. Repositories such as [`dreamer4`](https://github.com/lucidrains/dreamer4), [`metacontroller`](https://github.com/lucidrains/metacontroller), [`mimic-video`](https://github.com/lucidrains/mimic-video), [`pi-zero-pytorch`](https://github.com/lucidrains/pi-zero-pytorch), [`sdft-pytorch`](https://github.com/lucidrains/sdft-pytorch), and [`locoformer`](https://github.com/lucidrains/locoformer) repeatedly need operations such as `align_dims_left`, `shape_with_replace`, `lens_to_mask`, `pad_sequence`, `safe_cat`, and `pack_with_inverse`. This package centralizes those operations in one typed import surface instead of re-implementing the same tensor utility layer in each model repository.
 
-If you already know `torch-einops-utils`, `torch_einops_kit` began as a typed substitute for that package and has since grown into a superset. In addition to everything from `torch-einops-utils`, this repository centralizes small utility functions that appear repeatedly in other lucidrains model repositories but were never collected in one place, such as `l2norm`, `once`, `pack_one`, and `unpack_one`. The function family remains the same kind: small PyTorch and `einops` helpers for shape work, masks, padding, optional tensors, PyTree traversal, device routing, and checkpoint reconstruction. The import path is `torch_einops_kit`, not `torch_einops_utils`. The relationship is conceptual, not literal import-path compatibility.
+If you already know `torch-einops-utils`, `torch_einops_kit` began as a typed substitute for that package and has since grown into a superset. In addition to everything from `torch-einops-utils`, this repository centralizes small utility functions that appear repeatedly in other lucidrains model repositories but were never collected in one place, such as `l2norm`, `once`, `pack_one`, and `unpack_one`. The function family remains the same kind: small PyTorch and `einops` helpers for shape work, masks, padding, optional tensors, PyTree traversal, device routing, lightweight `nn.Module` adaptation, and checkpoint reconstruction. The import path is `torch_einops_kit`, not `torch_einops_utils`. The relationship is conceptual, not literal import-path compatibility.
 
 Use `torch_einops_kit` when you want strict typing, a `py.typed` marker, focused modules, extensive tests, and docstrings written for both humans and AI assistants. Use upstream when you want the most compact possible version of the same idea.
 
@@ -20,7 +20,7 @@ Use `torch_einops_kit` when you want strict typing, a `py.typed` marker, focused
 - Python requirement: `>=3.10`.
 - Runtime dependencies: `torch`, `einops`, and `typing-extensions`.
 - Root package exports: helper functions, slicing helpers, rank-alignment helpers, mask helpers, safe concatenation helpers, padding helpers, normalization helpers, and PyTree / `einops` helpers.
-- Submodules with dedicated imports: `torch_einops_kit.device`, `torch_einops_kit.einops`, `torch_einops_kit.save_load`, and `torch_einops_kit.scaleValues`.
+- Submodules with dedicated imports: `torch_einops_kit.device`, `torch_einops_kit.einops`, `torch_einops_kit.nn`, `torch_einops_kit.save_load`, and `torch_einops_kit.scaleValues`.
 - Typing status: the package ships a `py.typed` marker and the repository uses strict type checking.
 - Best fit: lucidrains-style model repositories that work with variable-length tensors, `einops` patterns, optional intermediate tensors, and nested `torch.nn.Module` graphs.
 
@@ -82,6 +82,16 @@ from torch_einops_kit.device import (
 )
 ```
 
+Import lightweight `nn.Module` adapters from `torch_einops_kit.nn`:
+
+```python
+from torch_einops_kit.nn import (
+    Identity,
+    Lambda,
+    Sequential,
+)
+```
+
 Import checkpoint decorators from `torch_einops_kit.save_load`:
 
 ```python
@@ -92,7 +102,7 @@ from torch_einops_kit.save_load import (
 )
 ```
 
-Import checkpoint decorators from `torch_einops_kit.scaleValues`:
+Import normalization and masked-reduction helpers from `torch_einops_kit.scaleValues`:
 
 ```python
 from torch_einops_kit.scaleValues import (
@@ -262,10 +272,10 @@ When `pad_lens=True` and `return_lens=True`, the second tensor contains padding
 
 ### PyTree helpers
 
-| Name                              | Contract                                                                                                                                                               |
-| --------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `tree_map_tensor(fn, tree)`       | Applies `fn` to every tensor leaf in a PyTree and leaves non-tensor leaves unchanged.                                                                                  |
-| `tree_flatten_with_inverse(tree)` | Returns a flat list of leaves and an inverse function that reconstructs the original PyTree shape from a replacement iterable of leaves.                               |
+| Name                              | Contract                                                                                                                                 |
+| --------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
+| `tree_map_tensor(fn, tree)`       | Applies `fn` to every tensor leaf in a PyTree and leaves non-tensor leaves unchanged.                                                    |
+| `tree_flatten_with_inverse(tree)` | Returns a flat list of leaves and an inverse function that reconstructs the original PyTree shape from a replacement iterable of leaves. |
 
 ## `scaleValues` submodule reference
 
@@ -298,6 +308,16 @@ The `torch_einops_kit.device` submodule contains three utilities for device infe
 | `move_inputs_to_device(device)`    | Decorator that recursively moves every tensor inside positional and keyword arguments to `device` before calling the wrapped function. Non-tensor values pass through unchanged.                                                                                |
 | `move_inputs_to_module_device(fn)` | 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. |
 
+## `nn` submodule reference
+
+The `torch_einops_kit.nn` submodule contains lightweight `torch.nn.Module` adapters and a `Sequential` constructor that ignores `None` values.
+
+| Name                   | Contract                                                                                                                                                                                                       |
+| ---------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `Sequential(*modules)` | Equivalent to `nn.Sequential(*compact(modules))`. Every `None` argument is discarded before construction.                                                                                                      |
+| `Identity()`           | `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. |
+| `Lambda(fn)`           | `torch.nn.Module` wrapper that stores `fn` on `self.fn` and delegates `forward(*args, **kwargs)` to `fn(*args, **kwargs)` without changing the argument structure.                                             |
+
 ## `save_load` submodule reference
 
 The `torch_einops_kit.save_load` submodule contains the checkpoint decorator and the two advanced configuration helpers that support nested decorated modules.
@@ -372,7 +392,7 @@ If you are an AI assistant adapting code from `torch-einops-utils`, use these tr
 This repository is not a repackaged mirror of upstream. This repository makes a different trade-off.
 
 - Upstream is intentionally compact.
-- This fork splits the implementation across focused modules such as `_helpers.py`, `_padding.py`, `device.py`, and `save_load.py` while still re-exporting most tensor helpers from the package root.
+- This fork splits the implementation across focused modules such as `_helpers.py`, `_padding.py`, `device.py`, `nn.py`, and `save_load.py` while still re-exporting most tensor helpers from the package root.
 - This fork adds strict typing, a `py.typed` marker, extensive tests, and detailed docstrings.
 - This fork is best treated as a typed, documented branch of the same utility idea rather than a literal import-path-compatible drop-in replacement.
 
@@ -380,6 +400,7 @@ This repository is not a repackaged mirror of upstream. This repository makes a
 
 - `src/torch_einops_kit/` — package source.
 - `src/torch_einops_kit/device.py` — device inference and input-routing decorators.
+- `src/torch_einops_kit/nn.py` — lightweight `torch.nn.Module` adapters and a `Sequential` constructor that ignores `None` modules.
 - `src/torch_einops_kit/save_load.py` — checkpoint save / load decorator and nested reconstruction helpers.
 - `src/torch_einops_kit/scaleValues.py` — vector normalization, masked mean, and the `RMSNorm` layer.
 - `tests/` — regression tests and usage examples for helpers, masks, padding, device routing, and checkpoint reconstruction.
@@ -401,7 +422,6 @@ pytest
 Run static analysis:
 
 ```bash
-pyright
 ruff check .
 ```
 

{torch_einops_kit-0.1.2 → torch_einops_kit-0.1.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "torch-einops-kit"
-version = "0.1.2"
+version = "0.1.3"
 description = "Typed tensor-shaping, masking, padding, device-routing, and checkpoint utilities for PyTorch and `einops`. A superset of `lucidrains/torch-einops-utils` with similar utilities from other lucidrains packages. Adds strict typing, extensive tests, and comprehensive docstrings."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -33,12 +33,12 @@ classifiers = [
 ]
 dependencies = [
   "einops>=0.8.2",
-  "torch>=2.10.0",
+  "torch>=2.0.0",
   "typing-extensions>=4.0.0",
 ]
 
 [project.urls]
-Documentation = "https://context7.com/hunterhogan/torch_einops_kit"
+Context7 = "https://context7.com/hunterhogan/torch_einops_kit"
 Donate = "https://www.patreon.com/integrated"
 Download = "https://pypi.org/project/torch_einops_kit"
 Homepage = "https://github.com/hunterhogan/torch_einops_kit"

{torch_einops_kit-0.1.2 → torch_einops_kit-0.1.3}/src/torch_einops_kit/__init__.py

@@ -1,8 +1,8 @@
-"""Access PyTorch tensor-shaping, masking, padding, and checkpoint utilities.
+"""Access PyTorch tensor-shaping, masking, padding, `nn.Module` adaptation, and checkpoint utilities.
 
 You can use this package for optional-value handling, tensor slicing, rank alignment, mask
-construction, safe concatenation, sequence padding, PyTree traversal, and `einops`-pattern tensor
-packing.
+construction, safe concatenation, sequence padding, PyTree traversal, lightweight `nn.Module`
+adapters, and `einops` pattern tensor packing.
 
 Helpers
 -------
@@ -103,16 +103,21 @@ device
 	Determine `torch.nn.Module` devices and decorate callables to move `Tensor` arguments automatically.
 einops
 	Pack and unpack `Tensor` objects with `einops` patterns and paired inverse functions.
+nn
+	Adapt callables and optional module sequences to the `torch.nn.Module` interface.
 save_load
 	Decorate `torch.nn.Module` subclasses with checkpoint save, load, and reconstruction helpers.
 scaleValues
 	Compute exclusive prefix sums, normalize feature vectors, and compute masked means.
 """
+
 # isort: split
+from __future__ import annotations
+
 from torch_einops_kit._semiotics import decreasing as decreasing, zeroIndexed as zeroIndexed
 
 # isort: split
-from torch_einops_kit._types import (
+from torch_einops_kit._theTypes import (
 	ConfigArgsKwargs as ConfigArgsKwargs, DehydratedCheckpoint as DehydratedCheckpoint, DehydratedTorchNNModule as DehydratedTorchNNModule,
 	DimAndValue as DimAndValue, IdentityCallable as IdentityCallable, PSpec as PSpec, RVar as RVar, StrPath as StrPath,
 	SupportsIntIndex as SupportsIntIndex, T_co as T_co, TorchNNModule as TorchNNModule, TVar as TVar)
@@ -151,4 +156,8 @@ from torch_einops_kit.utils import tree_flatten_with_inverse as tree_flatten_wit
 # isort: split
 # NOTE These imports are for backwards compatibility. Linters ought to tell users to import from the correct submodules.
 from torch_einops_kit.einops import pack_with_inverse  # pyright: ignore[reportUnusedImport]
+from torch_einops_kit.nn import Identity, Lambda, Sequential  # pyright: ignore[reportUnusedImport]
 from torch_einops_kit.scaleValues import l2norm, masked_mean  # pyright: ignore[reportUnusedImport]
+
+# NOTE `broadcat` is the identifier used in lucidrains packages.
+broadcat = broadcast_cat

{torch_einops_kit-0.1.2 → torch_einops_kit-0.1.3}/src/torch_einops_kit/_cat_and_stack.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
 from torch import broadcast_tensors, cat, stack, Tensor  # pyright: ignore[reportUnknownVariableType]
 from torch_einops_kit import safe
-from typing import cast
+from typing import cast, TYPE_CHECKING
+
+if TYPE_CHECKING:
+	from collections.abc import Sequence
 
 def broadcast_cat(tensors: Sequence[Tensor], dim: int = -1) -> Tensor:
 	"""Broadcast tensor groups before concatenation.
@@ -29,10 +31,10 @@ def broadcast_cat(tensors: Sequence[Tensor], dim: int = -1) -> Tensor:
 	[1] PyTorch - Context7
 		https://context7.com/pytorch/pytorch
 	"""
-	return cat(cast(list[Tensor], broadcast_tensors(*tensors)), dim)
+	return cat(cast('list[Tensor]', broadcast_tensors(*tensors)), dim)
 
 @safe
-def safe_cat(tensors: Sequence[Tensor], dim: int = 0) -> Tensor | None:
+def safe_cat(tensors: tuple[Tensor, ...] | list[Tensor], dim: int = 0) -> Tensor | None:
 	"""Concatenate tensors from `tensors` along an existing dimension, skipping `None` values.
 
 	You can use `safe_cat` to concatenate a mixed sequence of `Tensor` and `None` values. The `safe`
@@ -45,7 +47,7 @@ def safe_cat(tensors: Sequence[Tensor], dim: int = 0) -> Tensor | None:
 
 	Parameters
 	----------
-	tensors : Sequence[Tensor | None]
+	tensors : tuple[Tensor | None, ...] | list[Tensor | None]
 		A sequence of `Tensor` or `None` values. `None` values are filtered out before concatenation.
 		All non-`None` `Tensor` values must have the same shape in every dimension except `dim`.
 	dim : int = 0
@@ -79,9 +81,9 @@ def safe_cat(tensors: Sequence[Tensor], dim: int = 0) -> Tensor | None:
 	From sdft_pytorch [4], accumulating per-step token losses across a generation loop where
 	`token_kl_div_losses` is initialized to `None` before the loop:
 
-		```python
+	```python
 		token_kl_div_losses = safe_cat((token_kl_div_losses, token_kl_div), dim=1)
-		```
+	```
 
 	References
 	----------
@@ -94,10 +96,10 @@ def safe_cat(tensors: Sequence[Tensor], dim: int = 0) -> Tensor | None:
 	[4] lucidrains/sdft-pytorch
 		https://github.com/lucidrains/sdft-pytorch
 	"""
-	return cat(tensors, dim = dim)  # pyright: ignore[reportUnknownVariableType, reportCallIssue, reportArgumentType] https://github.com/pytorch/pytorch/issues/179391  # ty:ignore[no-matching-overload]
+	return None if len(tensors) == 0 else cat(tensors, dim=dim)
 
 @safe
-def safe_stack(tensors: Sequence[Tensor], dim: int = 0) -> Tensor | None:
+def safe_stack(tensors: tuple[Tensor, ...] | list[Tensor], dim: int = 0) -> Tensor | None:
 	"""Stack tensors from `tensors` along a new dimension, skipping `None` values.
 
 	You can use `safe_stack` to stack a mixed sequence of `Tensor` and `None` values. The `safe` [1]
@@ -106,8 +108,8 @@ def safe_stack(tensors: Sequence[Tensor], dim: int = 0) -> Tensor | None:
 
 	Parameters
 	----------
-	tensors : Sequence[Tensor | None]
-		A `Sequence` of `Tensor` or `None` values. `None` values are filtered out before stacking.
+	tensors : tuple[Tensor | None, ...] | list[Tensor | None]
+		A sequence of `Tensor` or `None` values. `None` values are filtered out before stacking.
 		All non-`None` `Tensor` values must have the same shape.
 	dim : int = 0
 		The dimension along which to stack. The result has one more dimension than each input
@@ -163,10 +165,10 @@ def safe_stack(tensors: Sequence[Tensor], dim: int = 0) -> Tensor | None:
 	[4] lucidrains/dreamer4
 		https://github.com/lucidrains/dreamer4
 	"""
-	return stack(tensors, dim = dim)  # pyright: ignore[reportArgumentType] https://github.com/pytorch/pytorch/issues/179391  # ty:ignore[invalid-argument-type]
+	return None if len(tensors) == 0 else stack(tensors, dim=dim)
 
 """
-Some or all of the logic in this module may be protected by the following.
+Some of the logic in this module may be protected by the following.
 
 MIT License
 

{torch_einops_kit-0.1.2 → torch_einops_kit-0.1.3}/src/torch_einops_kit/_dimensions.py

@@ -1,6 +1,11 @@
-from collections.abc import Sequence
-from torch import Tensor
+from __future__ import annotations
+
 from torch_einops_kit import exists
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+	from collections.abc import Sequence
+	from torch import Tensor
 
 def pad_right_ndim(t: Tensor, ndims: int) -> Tensor:
 	"""Reshape a tensor by inserting singleton dimensions at the trailing side.
@@ -45,11 +50,7 @@ def pad_right_ndim(t: Tensor, ndims: int) -> Tensor:
 	"""
 	return pad_ndim(t, (0, ndims))
 
-def align_dims_left(
-	tensors: Sequence[Tensor],
-	*,
-	ndim: int | None = None,
-) -> tuple[Tensor, ...]:
+def align_dims_left(tensors: Sequence[Tensor], *, ndim: int | None = None) -> tuple[Tensor, ...]:
 	"""Pad all tensors in a sequence with trailing singleton dimensions to a common rank.
 
 	You can use this function to align a heterogeneous sequence of tensors to the same number of
@@ -87,33 +88,33 @@ def align_dims_left(
 	Align a PPO advantage tensor `(b, n)` with a log-probability ratio tensor `(b, n, d)` for
 	element-wise multiplication:
 
-		```python
+	```python
 		from torch_einops_kit import align_dims_left
 
 		# metacontroller: align ratio and advantages before the PPO surrogate loss
 		ratio, advantages = align_dims_left((ratio, advantages))
 		surr1 = ratio * advantages
-		```
+	```
 
 	Align a noise schedule `(b,)` with a latent tensor `(b, n, d)` for linear interpolation:
 
-		```python
+	```python
 		from torch_einops_kit import align_dims_left
 
 		# dreamer4: align time with latents before noising
 		aligned_times, _ = align_dims_left((times, latents))
 		noised_latents = noise.lerp(latents, aligned_times)
-		```
+	```
 
 	Align a 1-D time value with an action tensor before flow-matching noise interpolation:
 
-		```python
+	```python
 		from torch_einops_kit import align_dims_left
 
 		# mimic_video: align time with actions for noise interpolation
 		actions, left_aligned_time = align_dims_left((actions, time))
 		noised = noise.lerp(actions, left_aligned_time)
-		```
+	```
 
 	References
 	----------
@@ -179,7 +180,7 @@ def pad_ndim(t: Tensor, ndims: tuple[int, int]) -> Tensor:
 	shape: tuple[int, ...] = t.shape
 	left, right = ndims
 	if left < 0 or right < 0:
-		message: str = f"I received `{left = }` and `{right = }`, but I need both values to be greater than or equal to `0`."
+		message: str = f'I received `{left = }` and `{right = }`, but I need both values to be greater than or equal to `0`.'
 		raise ValueError(message)
 
 	ones: tuple[int] = (1,)
@@ -311,22 +312,22 @@ def pad_right_ndim_to(t: Tensor, ndims: int) -> Tensor:
 	Broadcast a scalar time value against a video tensor of shape `(b, c, t, h, w)`
 	for flow interpolation:
 
-		```python
+	```python
 		from torch_einops_kit import pad_right_ndim_to
 
 		# dreamer4: align time '(b,)' with video '(b, c, t, h, w)'
 		padded_time = pad_right_ndim_to(time[None], video.ndim)
 		pred_flow = (pred_video - video) / (1.0 - padded_time)
-		```
+	```
 
 	Scale a flow prediction using a denominator with lower rank than the prediction:
 
-		```python
+	```python
 		from torch_einops_kit import pad_right_ndim_to
 
 		# mimic_video: convert model output to flow space
 		pred_flow = (pred - actions) / pad_right_ndim_to(1.0 - action_time, pred.ndim).clamp_min(eps)
-		```
+	```
 
 	References
 	----------