torch-einops-utils 0.0.9__tar.gz

torch_einops_utils-0.0.9/.github/workflows/python-publish.yml

@@ -0,0 +1,36 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build
+    - name: Build package
+      run: python -m build
+    - name: Publish package
+      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      with:
+        user: __token__
+        password: ${{ secrets.PYPI_API_TOKEN }}

torch_einops_utils-0.0.9/.github/workflows/test.yml

@@ -0,0 +1,18 @@
+name: Pytest
+on: [push, pull_request]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: true
+        python-version: "3.10"
+
+    - name: Test with pytest
+      run: uv run --extra test pytest tests/

torch_einops_utils-0.0.9/.gitignore

@@ -0,0 +1,207 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/

torch_einops_utils-0.0.9/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Phil Wang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torch_einops_utils-0.0.9/PKG-INFO

@@ -0,0 +1,46 @@
+Metadata-Version: 2.4
+Name: torch-einops-utils
+Version: 0.0.9
+Summary: Personal utility functions
+Project-URL: Homepage, https://pypi.org/project/torch-einops-utils/
+Project-URL: Repository, https://github.com/lucidrains/torch-einops-utils
+Author-email: Phil Wang <lucidrains@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2026 Phil Wang
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: einops,torch
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: einops>=0.8.1
+Requires-Dist: torch>=2.5
+Provides-Extra: examples
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Description-Content-Type: text/markdown
+
+## torch-einops-utils
+
+Some utility functions to help myself (and perhaps others) go faster with ML/AI work

torch_einops_utils-0.0.9/README.md

@@ -0,0 +1,3 @@
+## torch-einops-utils
+
+Some utility functions to help myself (and perhaps others) go faster with ML/AI work

torch_einops_utils-0.0.9/pyproject.toml

@@ -0,0 +1,56 @@
+[project]
+name = "torch-einops-utils"
+version = "0.0.9"
+description = "Personal utility functions"
+authors = [
+    { name = "Phil Wang", email = "lucidrains@gmail.com" }
+]
+readme = "README.md"
+requires-python = ">= 3.9"
+license = { file = "LICENSE" }
+keywords = [
+    'torch',
+    'einops'
+]
+
+classifiers=[
+    'Development Status :: 4 - Beta',
+    'Intended Audience :: Developers',
+    'Topic :: Scientific/Engineering :: Artificial Intelligence',
+    'License :: OSI Approved :: MIT License',
+    'Programming Language :: Python :: 3.9',
+]
+
+dependencies = [
+    "torch>=2.5",
+    "einops>=0.8.1"
+]
+
+[project.urls]
+Homepage = "https://pypi.org/project/torch-einops-utils/"
+Repository = "https://github.com/lucidrains/torch-einops-utils"
+
+[project.optional-dependencies]
+examples = []
+test = [
+    "pytest"
+]
+
+[tool.pytest.ini_options]
+pythonpath = [
+  "."
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.rye]
+managed = true
+dev-dependencies = []
+
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.build.targets.wheel]
+packages = ["torch_einops_utils"]

torch_einops_utils-0.0.9/tests/test_utils.py

@@ -0,0 +1,94 @@
+import torch
+from torch import tensor
+
+from torch_einops_utils.torch_einops_utils import (
+    exists,
+    pad_ndim,
+    align_dims_left,
+    pad_at_dim,
+    pad_left_at_dim,
+    pad_right_at_dim,
+    pad_sequence,
+    lens_to_mask,
+    and_masks,
+    or_masks,
+    tree_flatten_with_inverse,
+    pack_with_inverse,
+)
+
+def test_exist():
+    assert not exists(None)
+
+def test_pad_ndim():
+    t = torch.randn(3)
+    t = pad_ndim(t, (1, 2))
+    assert t.shape == (1, 3, 1, 1)
+
+def test_align_ndim_left():
+    t = torch.randn(3)
+    u = torch.randn(3, 5, 2)
+    v = torch.randn(3, 5)
+
+    t, u, v = align_dims_left((t, u, v))
+    assert t.shape == (3, 1, 1)
+    assert u.shape == (3, 5, 2)
+    assert v.shape == (3, 5, 1)
+
+def test_pad_at_dim():
+    t = torch.randn(3, 6, 1)
+    padded = pad_at_dim(t, (0, 1), dim = 1)
+
+    assert padded.shape == (3, 7, 1)
+    assert torch.allclose(padded, pad_right_at_dim(t, 1, dim = 1))
+    assert not torch.allclose(padded, pad_left_at_dim(t, 1, dim = 1))
+
+def test_tree_flatten_with_inverse():
+    tree = (1, (2, 3), 4)
+    (first, *rest), inverse = tree_flatten_with_inverse(tree)
+
+    out = inverse((first + 1, *rest))
+    assert out == (2, (2, 3), 4)
+
+def test_pack_with_inverse():
+    t = torch.randn(3, 12, 2, 2)
+    t, inverse = pack_with_inverse(t, 'b * d')
+
+    assert t.shape == (3, 24, 2)
+    t = inverse(t)
+    assert t.shape == (3, 12, 2, 2)
+
+    u = torch.randn(3, 4, 2)
+    t, inverse = pack_with_inverse([t, u], 'b * d')
+    assert t.shape == (3, 28, 2)
+
+    t = t.sum(dim = -1)
+    t, u = inverse(t, 'b *')
+    assert t.shape == (3, 12, 2)
+    assert u.shape == (3, 4)
+
+def test_better_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)
+    assert packed.shape == (3, 2, 4, 5)
+    assert lens.tolist() == [4, 3, 1]
+
+    mask = lens_to_mask(lens)
+    assert torch.allclose(mask.sum(dim = -1), lens)
+
+def test_and_masks():
+    assert not exists(and_masks([None]))
+
+    mask1 = tensor([True, True])
+    mask2 = tensor([True, False])
+    assert (and_masks([mask1, None, mask2]) == tensor([True, False])).all()
+
+def test_or_masks():
+    assert not exists(or_masks([None]))
+
+    mask1 = tensor([True, True])
+    mask2 = tensor([True, False])
+    assert (or_masks([mask1, None, mask2]) == tensor([True, True])).all()

torch_einops_utils-0.0.9/torch_einops_utils/__init__.py

@@ -0,0 +1,26 @@
+
+from torch_einops_utils.torch_einops_utils import (
+    pad_ndim,
+    pad_left_ndim,
+    pad_right_ndim,
+    align_dims_left,
+)
+
+from torch_einops_utils.torch_einops_utils import (
+    lens_to_mask
+)
+
+from torch_einops_utils.torch_einops_utils import (
+    pad_at_dim,
+    pad_left_at_dim,
+    pad_right_at_dim,
+    pad_sequence
+)
+
+from torch_einops_utils.torch_einops_utils import (
+    tree_flatten_with_inverse
+)
+
+from torch_einops_utils.torch_einops_utils import (
+    pack_with_inverse
+)

torch_einops_utils-0.0.9/torch_einops_utils/torch_einops_utils.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+from typing import Literal
+
+import torch
+from torch import tensor, is_tensor, cat, stack, arange
+import torch.nn.functional as F
+
+from torch.utils._pytree import tree_flatten, tree_unflatten
+
+from einops import rearrange, repeat, reduce, pack, unpack
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def first(arr):
+    return arr[0]
+
+# exported functions
+
+# dimensions
+
+def pad_ndim(t, ndims: tuple[int, int]):
+    shape = t.shape
+    left, right = ndims
+
+    ones = (1,)
+    ones_left = ones * left
+    ones_right = ones * right
+    return t.reshape(*ones_left, *shape, *ones_right)
+
+def pad_left_ndim(t, ndims: int):
+    return pad_ndim(t, (ndims, 0))
+
+def pad_right_ndim(t, ndims: int):
+    return pad_ndim(t, (0, ndims))
+
+def align_dims_left(
+    tensors,
+    *,
+    ndim = None
+):
+    if not exists(ndim):
+        ndim = max([t.ndim for t in tensors])
+
+    return tuple(pad_right_ndim(t, ndim - t.ndim) for t in tensors)
+
+# masking
+
+def lens_to_mask(lens, max_len = None):
+    device = lens.device
+
+    if not exists(max_len):
+        max_len = lens.amax().item()
+
+    seq = arange(max_len, device = device)
+    lens = rearrange(lens, '... -> ... 1')
+    return seq < lens
+
+def reduce_masks(masks, op):
+    masks = [*filter(exists, masks)]
+
+    if len(masks) == 0:
+        return None
+    elif len(masks) == 1:
+        return first(masks)
+
+    mask, *rest_masks = masks
+
+    for rest_mask in rest_masks:
+        mask = op(mask, rest_mask)
+
+    return mask
+
+def and_masks(masks):
+    return reduce_masks(masks, torch.logical_and)
+
+def or_masks(masks):
+    return reduce_masks(masks, torch.logical_or)
+
+# padding
+
+def pad_at_dim(
+    t,
+    pad: tuple[int, int],
+    *,
+    dim = -1,
+    value = 0.
+):
+    dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
+    zeros = ((0, 0) * dims_from_right)
+    return F.pad(t, (*zeros, *pad), value = value)
+
+def pad_left_at_dim(t, pad: int, **kwargs):
+    return pad_at_dim(t, (pad, 0), **kwargs)
+
+def pad_right_at_dim(t, pad: int, **kwargs):
+    return pad_at_dim(t, (0, pad), **kwargs)
+
+# better pad sequence
+
+def pad_sequence(
+    tensors,
+    *,
+    dim = -1,
+    value = 0.,
+    left = False,
+    dim_stack = 0,
+    return_lens = False
+):
+    if len(tensors) == 0:
+        return None
+
+    device = first(tensors).device
+
+    lens = tensor([t.shape[dim] for t in tensors], device = device)
+    max_len = lens.amax().item()
+
+    pad_fn = pad_left_at_dim if left else pad_right_at_dim
+    padded_tensors = [pad_fn(t, max_len - t_len, dim = dim, value = value) for t, t_len in zip(tensors, lens)]
+
+    stacked = stack(padded_tensors, dim = dim_stack)
+
+    if not return_lens:
+        return stacked
+
+    return stacked, lens
+
+# tree flatten with inverse
+
+def tree_flatten_with_inverse(tree):
+    flattened, spec = tree_flatten(tree)
+
+    def inverse(out):
+        return tree_unflatten(out, spec)
+
+    return flattened, inverse
+
+# einops pack
+
+def pack_with_inverse(t, pattern):
+    is_one = is_tensor(t)
+
+    if is_one:
+        t = [t]
+
+    packed, packed_shape = pack(t, pattern)
+
+    def inverse(out, inv_pattern = None):
+        inv_pattern = default(inv_pattern, pattern)
+        out = unpack(out, packed_shape, inv_pattern)
+
+        if is_one:
+            out = first(out)
+
+        return out
+
+    return packed, inverse