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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-einops-kit
-Version: 0.1.3
+Version: 0.1.4
 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

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

@@ -1,6 +1,6 @@
 [project]
 name = "torch-einops-kit"
-version = "0.1.3"
+version = "0.1.4"
 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"

{torch_einops_kit-0.1.3 → torch_einops_kit-0.1.4}/src/torch_einops_kit/_padding.py

@@ -289,32 +289,42 @@ def pad_sequence(
 	return_stacked: Literal[True] = True,
 	return_lens: Literal[False] = False,
 	pad_lens: bool = False
-) -> Tensor | None: ...
-# `list[Tensor] | None` could create confusion: see the tuple | None overloads.
+) -> Tensor: ...
 @overload
 def pad_sequence(
 	tensors: Sequence[Tensor], *, dim: int = -1, value: float = 0., left: bool = False, dim_stack: int = 0,
 	return_stacked: Literal[False],
 	return_lens: Literal[False] = False,
 	pad_lens: bool = False
-) -> list[Tensor] | None: ...
-# TODO return `tuple | None` is problematic. See `test_better_pad_sequence`. Pyright gives a confusing error: "None is not iterable" when trying to unpack the tuple.
+) -> list[Tensor]: ...
 @overload
 def pad_sequence(
 	tensors: Sequence[Tensor], *, dim: int = -1, value: float = 0., left: bool = False, dim_stack: int = 0,
 	return_stacked: Literal[True] = True,
 	return_lens: Literal[True],
 	pad_lens: bool = False
-) -> tuple[Tensor, Tensor] | None: ...
+) -> tuple[Tensor, Tensor]: ...
 @overload
 def pad_sequence(
 	tensors: Sequence[Tensor], *, dim: int = -1, value: float = 0., left: bool = False, dim_stack: int = 0,
 	return_stacked: Literal[False],
 	return_lens: Literal[True],
 	pad_lens: bool = False
-) -> tuple[list[Tensor], Tensor] | None: ...
+) -> tuple[list[Tensor], Tensor]: ...
+@overload
+def pad_sequence(
+	tensors: list[Tensor] = [],  # noqa: B006
+	*, dim: int = -1, value: float = 0.0, left: bool = False, dim_stack: int = 0,
+	return_stacked: bool = True, return_lens: bool = False, pad_lens: bool = False
+) -> None: ...
+@overload
+def pad_sequence(
+	tensors: tuple[()] = (),
+	*, dim: int = -1, value: float = 0.0, left: bool = False, dim_stack: int = 0,
+	return_stacked: bool = True, return_lens: bool = False, pad_lens: bool = False
+) -> None: ...
 def pad_sequence(
-	tensors: Sequence[Tensor],
+	tensors: Sequence[Tensor] = (),
 	*,
 	dim: int = -1,
 	value: float = 0.0,
@@ -443,7 +453,13 @@ def pad_sequence(
 
 	return output
 
-def pad_sequence_and_cat(tensors: Sequence[Tensor], *, dim_cat: int = 0, dim: int = -1, value: float = 0., left: bool = False) -> Tensor | None:
+@overload
+def pad_sequence_and_cat(tensors: Sequence[Tensor], *, dim_cat: int = 0, dim: int = -1, value: float = 0., left: bool = False) -> Tensor:...
+@overload
+def pad_sequence_and_cat(tensors: list[Tensor] = [], *, dim_cat: int = 0, dim: int = -1, value: float = 0., left: bool = False) -> None:...  # noqa: B006
+@overload
+def pad_sequence_and_cat(tensors: tuple[()]=(), *, dim_cat: int = 0, dim: int = -1, value: float = 0., left: bool = False) -> None:...
+def pad_sequence_and_cat(tensors: Sequence[Tensor]=(), *, dim_cat: int = 0, dim: int = -1, value: float = 0., left: bool = False) -> Tensor | None:
 	"""Pad `tensors` to a shared length along `dim` and concatenate along `dim_cat`.
 
 	You can use `pad_sequence_and_cat` to align and merge a heterogeneous-length sequence of `Tensor`
@@ -498,12 +514,12 @@ def pad_sequence_and_cat(tensors: Sequence[Tensor], *, dim_cat: int = 0, dim: in
 	[3] tests/test_utils.py
 	"""
 	padded: Tensor | list[Tensor] | None = pad_sequence(tensors, dim=dim, value=value, left=left, return_stacked=False, return_lens=False)
-	if padded is not None:
+	if padded is not None: # pyright: ignore[reportUnnecessaryComparison]
 		return cat(padded, dim=dim_cat)
 	return padded
 
 """
-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.3 → torch_einops_kit-0.1.4}/src/torch_einops_kit/scaleValues.py

@@ -107,7 +107,8 @@ def l2norm(t: Tensor) -> Tensor:
 	"""
 	return F.normalize(t, dim=-1, p=2)
 
-def masked_mean(t: Tensor, mask: Tensor | None = None, dim: torch.Size | list[int] | tuple[int, ...] | int | None = None, eps: float = 1e-5) -> Tensor:
+def masked_mean(t: Tensor, mask: Tensor | None = None, dim: torch.Size | list[int] | tuple[int, ...] | int | None = None, eps: float = 1e-5
+				, *, keepdim: bool = False) -> Tensor:
 	"""Compute the mean of `t` over positions selected by `mask`.
 
 	You can use this function to average only the elements of `t` where `mask` is `True`, ignoring
@@ -131,6 +132,9 @@ def masked_mean(t: Tensor, mask: Tensor | None = None, dim: torch.Size | list[in
 	eps : float = 1e-5
 		A small value added to the denominator to prevent division by zero when computing the masked
 		mean along a dimension.
+	keepdim : bool = False
+		When `True`, the reduced dimension is retained as a size-one dimension in `result`, allowing
+		`result` to broadcast against `t`.
 
 	Returns
 	-------
@@ -172,6 +176,15 @@ def masked_mean(t: Tensor, mask: Tensor | None = None, dim: torch.Size | list[in
 		# result == tensor([1.0, 3.5])
 	```
 
+	Compute the sequence-level mean of transformer activations with `keepdim=True` to preserve the
+	sequence dimension for broadcasting, used in the mean-variance split residual update [4][5]:
+
+	```python
+		mean_x = masked_mean(x, mask=mask, dim=-2, keepdim=True)
+		mean_residual = masked_mean(residual, mask=mask, dim=-2, keepdim=True)
+		centered_x = x - mean_x
+	```
+
 	References
 	----------
 	[1] torch.Tensor.mean - PyTorch documentation
@@ -180,20 +193,20 @@ def masked_mean(t: Tensor, mask: Tensor | None = None, dim: torch.Size | list[in
 
 	[3] tests.test_utils.test_masked_mean
 
+	[4] Lu, P. (2026). Mean Mode Screaming: Mean-Variance Split Residuals for 1000-Layer Diffusion
+		Transformers. arXiv:2605.06169.
+		https://arxiv.org/abs/2605.06169
+	[5] x_transformers.x_transformers.MVSplitResidualUpdate.forward - lucidrains/x-transformers
+		https://github.com/lucidrains/x-transformers
+
 	"""
 	if not exists(mask):
-		return t.mean(dim=dim) if exists(dim) else t.mean()
-
-	if mask.ndim < t.ndim:
-		mask = pad_right_ndim(mask, t.ndim - mask.ndim)
-
-	mask = mask.expand_as(t)
+		mask = torch.ones_like(t, dtype=torch.bool)
 
-	if not exists(dim):
-		return t[mask].mean() if mask.any() else t[mask].sum()
+	mask = pad_right_ndim(mask, max(0, t.ndim - mask.ndim)).expand_as(t)
 
-	num: Tensor = (t * mask).sum(dim=dim)
-	den: Tensor = mask.sum(dim=dim)
+	num: Tensor = (t * mask).sum(dim=dim, keepdim=keepdim)
+	den: Tensor = mask.sum(dim=dim, keepdim=keepdim)
 
 	return num / den.clamp(min=eps)