torch-max-mem 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

torch_max_mem/__init__.py

@@ -1,9 +1,8 @@
-# -*- coding: utf-8 -*-
-
 """Maximize memory utilization with PyTorch."""
+
 from .api import MemoryUtilizationMaximizer, maximize_memory_utilization
 
 __all__ = [
-    "maximize_memory_utilization",
     "MemoryUtilizationMaximizer",
+    "maximize_memory_utilization",
 ]

torch_max_mem/api.py

@@ -1,5 +1,3 @@
-# -*- coding: utf-8 -*-
-
 """
 This module contains the public API.
 
@@ -9,6 +7,7 @@ Assume you have a function for batched computation of nearest neighbors using br
 
     import torch
 
+
     def knn(x, y, batch_size, k: int = 3):
         return torch.cat(
             [
@@ -26,6 +25,7 @@ out-of-memory error occurs.
     import torch
     from torch_max_mem import maximize_memory_utilization
 
+
     @maximize_memory_utilization()
     def knn(x, y, batch_size, k: int = 3):
         return torch.cat(
@@ -45,6 +45,7 @@ In the code, you can now always pass the largest sensible batch size, e.g.,
     y = torch.rand(200, 100, device="cuda")
     knn(x, y, batch_size=x.shape[0])
 """
+
 # cf. https://gist.github.com/mberr/c37a8068b38cabc98228db2cbe358043
 from __future__ import annotations
 
@@ -52,17 +53,10 @@ import functools
 import inspect
 import itertools
 import logging
-from collections import ChainMap
+from collections.abc import Collection, Iterable, Mapping, MutableMapping, Sequence
 from typing import (
     Any,
     Callable,
-    Collection,
-    Iterable,
-    Mapping,
-    MutableMapping,
-    Optional,
-    Sequence,
-    Tuple,
     TypeVar,
 )
 
@@ -99,9 +93,7 @@ def upgrade_to_sequence(
         when the (inferred) length of q and parameter_name do not match
     """
     # normalize parameter name
-    parameter_names = (
-        (parameter_name,) if isinstance(parameter_name, str) else tuple(parameter_name)
-    )
+    parameter_names = (parameter_name,) if isinstance(parameter_name, str) else tuple(parameter_name)
     q = (q,) if isinstance(q, int) else tuple(q)
     q = q * len(parameter_names) if len(q) == 1 else q
     if len(q) != len(parameter_names):
@@ -128,7 +120,7 @@ def determine_default_max_value(
     :raises ValueError:
         when the function does not have a parameter of the given name
     """
-    if parameter_name not in signature.parameters.keys():
+    if parameter_name not in signature.parameters:
         raise ValueError(f"{func} does not have a parameter {parameter_name}.")
     _parameter = signature.parameters[parameter_name]
     if _parameter.annotation != inspect.Parameter.empty and _parameter.annotation not in (
@@ -146,10 +138,10 @@ def determine_default_max_value(
 
 def determine_max_value(
     bound_arguments: inspect.BoundArguments,
-    args: P.args,
-    kwargs: P.kwargs,
     parameter_name: str,
     default_max_value: int | Callable[P, int] | None,
+    *args: P.args,
+    **kwargs: P.kwargs,
 ) -> int:
     """
     Either use the provided value, or the default maximum value.
@@ -204,12 +196,13 @@ ADDITIONAL_OOM_ERROR_INFIXES = {
 def iter_tensor_devices(*args: Any, **kwargs: Any) -> Iterable[torch.device]:
     """Iterate over tensors' devices (may contain duplicates)."""
     for obj in itertools.chain(args, kwargs.values()):
-        if torch.is_tensor(obj):
-            assert isinstance(obj, torch.Tensor)
+        if isinstance(obj, torch.Tensor):
             yield obj.device
 
 
-def create_tensor_checker(safe_devices: Collection[str] | None = None) -> Callable[P, None]:
+def create_tensor_checker(
+    safe_devices: Collection[str] | None = None,
+) -> Callable[P, None]:
     """
     Create a function that warns when tensors are on any device that is not considered safe.
 
@@ -274,16 +267,15 @@ def is_oom_error(error: BaseException) -> bool:
         return True
     if not isinstance(error, RuntimeError):
         return False
-    if not error.args:
-        return False
-    return any(infix in error.args[0] for infix in ADDITIONAL_OOM_ERROR_INFIXES)
+    message = str(error)
+    return any(infix in message for infix in ADDITIONAL_OOM_ERROR_INFIXES)
 
 
 def maximize_memory_utilization_decorator(
     parameter_name: str | Sequence[str] = "batch_size",
     q: int | Sequence[int] = 32,
     safe_devices: Collection[str] | None = None,
-) -> Callable[[Callable[P, R]], Callable[P, Tuple[R, tuple[int, ...]]]]:
+) -> Callable[[Callable[P, R]], Callable[P, tuple[R, tuple[int, ...]]]]:
     """
     Create decorators to create methods for memory utilization maximization.
 
@@ -301,8 +293,8 @@ def maximize_memory_utilization_decorator(
     parameter_names, qs = upgrade_to_sequence(parameter_name, q)
 
     def decorator_maximize_memory_utilization(
-        func: Callable[P, R]
-    ) -> Callable[P, Tuple[R, tuple[int, ...]]]:
+        func: Callable[P, R],
+    ) -> Callable[P, tuple[R, tuple[int, ...]]]:
         """
         Decorate a function to maximize memory utilization.
 
@@ -320,9 +312,7 @@ def maximize_memory_utilization_decorator(
         }
 
         @functools.wraps(func)
-        def wrapper_maximize_memory_utilization(
-            *args: P.args, **kwargs: P.kwargs
-        ) -> Tuple[R, tuple[int, ...]]:
+        def wrapper_maximize_memory_utilization(*args: P.args, **kwargs: P.kwargs) -> tuple[R, tuple[int, ...]]:
             """
             Wrap a function to maximize memory utilization by successive halving.
 
@@ -345,11 +335,11 @@ def maximize_memory_utilization_decorator(
             # determine actual max values
             max_values = [
                 determine_max_value(
-                    bound_arguments=bound_arguments,
-                    args=args,
-                    kwargs=kwargs,
-                    parameter_name=name,
-                    default_max_value=default_max_value,
+                    bound_arguments,
+                    name,
+                    default_max_value,
+                    *args,
+                    **kwargs,
                 )
                 for name, default_max_value in default_max_values.items()
             ]
@@ -360,13 +350,11 @@ def maximize_memory_utilization_decorator(
 
             while i < len(max_values):
                 while max_values[i] > 0:
-                    p_kwargs = {
-                        name: max_value for name, max_value in zip(parameter_names, max_values)
-                    }
-                    # note: bound_arguments.kwargs is typed as dict, but (silently) immutable (=ignoring updates)...
-                    combined_kwargs: P.kwargs = ChainMap(p_kwargs, bound_arguments.kwargs)
+                    p_kwargs = dict(zip(parameter_names, max_values))
+                    # note: changes to arguments apply to both, .args and .kwargs
+                    bound_arguments.arguments.update(p_kwargs)
                     try:
-                        return func(*bound_arguments.args, **combined_kwargs), tuple(max_values)
+                        return func(*bound_arguments.args, **bound_arguments.kwargs), tuple(max_values)
                     except (torch.cuda.OutOfMemoryError, RuntimeError) as error:
                         # raise errors unrelated to out-of-memory
                         if not is_oom_error(error):
@@ -391,12 +379,8 @@ def maximize_memory_utilization_decorator(
                 i += 1
             # log memory summary for each CUDA device before raising memory error
             for device in {d for d in iter_tensor_devices(*args, **kwargs) if d.type == "cuda"}:
-                logger.debug(
-                    f"Memory summary for {device=}:\n{torch.cuda.memory_summary(device=device)}"
-                )
-            raise MemoryError(
-                f"Execution did not even succeed with {parameter_names} all equal to 1."
-            ) from last_error
+                logger.debug(f"Memory summary for {device=}:\n{torch.cuda.memory_summary(device=device)}")
+            raise MemoryError(f"Execution did not even succeed with {parameter_names} all equal to 1.") from last_error
 
         return wrapper_maximize_memory_utilization
 
@@ -455,7 +439,7 @@ class MemoryUtilizationMaximizer:
         parameter_name: str | Sequence[str] = "batch_size",
         q: int | Sequence[int] = 32,
         safe_devices: Collection[str] | None = None,
-        hasher: Optional[Callable[[Mapping[str, Any]], int]] = None,
+        hasher: Callable[[Mapping[str, Any]], int] | None = None,
         keys: Collection[str] | str | None = None,
     ) -> None:
         """
@@ -474,7 +458,7 @@ class MemoryUtilizationMaximizer:
         """
         self.parameter_names, self.qs = upgrade_to_sequence(parameter_name=parameter_name, q=q)
         self.safe_devices = safe_devices
-        self.parameter_value: MutableMapping[int, tuple[int, ...]] = dict()
+        self.parameter_value: MutableMapping[int, tuple[int, ...]] = {}
         if hasher is None:
             keys = KeyHasher.normalize_keys(keys)
             intersection = set(keys).intersection(self.parameter_names)

torch_max_mem/version.py

@@ -1,27 +1,25 @@
-# -*- coding: utf-8 -*-
-
 """Version information for :mod:`torch_max_mem`.
 
 Run with ``python -m torch_max_mem.version``
 """
 
 import os
-from subprocess import CalledProcessError, check_output  # noqa: S404
+from subprocess import CalledProcessError, check_output
 
 __all__ = [
     "VERSION",
-    "get_version",
     "get_git_hash",
+    "get_version",
 ]
 
-VERSION = "0.1.2"
+VERSION = "0.1.4"
 
 
 def get_git_hash() -> str:
     """Get the :mod:`torch_max_mem` git hash."""
     with open(os.devnull, "w") as devnull:
         try:
-            ret = check_output(  # noqa: S603,S607
+            ret = check_output(
                 ["git", "rev-parse", "HEAD"],
                 cwd=os.path.dirname(__file__),
                 stderr=devnull,

torch_max_mem-0.1.4.dist-info/METADATA

@@ -0,0 +1,345 @@
+Metadata-Version: 2.4
+Name: torch-max-mem
+Version: 0.1.4
+Summary: Maximize memory utilization with PyTorch.
+Keywords: snekpack,cookiecutter,torch
+Author: Max Berrendorf
+Author-email: Max Berrendorf <max.berrendorf@gmail.com>
+License-File: LICENSE
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Framework :: Pytest
+Classifier: Framework :: tox
+Classifier: Framework :: Sphinx
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3 :: Only
+Requires-Dist: torch>=2.0
+Requires-Dist: typing-extensions
+Maintainer: Max Berrendorf
+Maintainer-email: Max Berrendorf <max.berrendorf@gmail.com>
+Requires-Python: >=3.9
+Project-URL: Bug Tracker, https://github.com/mberr/torch-max-mem/issues
+Project-URL: Download, https://github.com/mberr/torch-max-mem/releases
+Project-URL: Homepage, https://github.com/mberr/torch-max-mem
+Project-URL: Source Code, https://github.com/mberr/torch-max-mem
+Description-Content-Type: text/markdown
+
+<!--
+<p align="center">
+  <img src="https://github.com/mberr/torch-max-mem/raw/main/docs/source/logo.png" height="150">
+</p>
+-->
+
+<h1 align="center">
+  torch-max-mem
+</h1>
+
+<p align="center">
+    <a href="https://github.com/mberr/torch-max-mem/actions/workflows/tests.yml">
+        <img alt="Tests" src="https://github.com/mberr/torch-max-mem/actions/workflows/tests.yml/badge.svg" /></a>
+    <a href="https://pypi.org/project/torch_max_mem">
+        <img alt="PyPI" src="https://img.shields.io/pypi/v/torch_max_mem" /></a>
+    <a href="https://pypi.org/project/torch_max_mem">
+        <img alt="PyPI - Python Version" src="https://img.shields.io/pypi/pyversions/torch_max_mem" /></a>
+    <a href="https://github.com/mberr/torch-max-mem/blob/main/LICENSE">
+        <img alt="PyPI - License" src="https://img.shields.io/pypi/l/torch_max_mem" /></a>
+    <a href='https://torch_max_mem.readthedocs.io/en/latest/?badge=latest'>
+        <img src='https://readthedocs.org/projects/torch_max_mem/badge/?version=latest' alt='Documentation Status' /></a>
+    <a href="https://codecov.io/gh/mberr/torch-max-mem/branch/main">
+        <img src="https://codecov.io/gh/mberr/torch-max-mem/branch/main/graph/badge.svg" alt="Codecov status" /></a>  
+    <a href="https://github.com/cthoyt/cookiecutter-python-package">
+        <img alt="Cookiecutter template from @cthoyt" src="https://img.shields.io/badge/Cookiecutter-snekpack-blue" /></a>
+    <a href="https://github.com/astral-sh/ruff">
+        <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json" alt="Ruff" style="max-width:100%;"></a>
+    <a href="https://github.com/mberr/torch-max-mem/blob/main/.github/CODE_OF_CONDUCT.md">
+        <img src="https://img.shields.io/badge/Contributor%20Covenant-2.1-4baaaa.svg" alt="Contributor Covenant"/></a>
+    <!-- uncomment if you archive on zenodo
+    <a href="https://zenodo.org/badge/latestdoi/XXXXXX">
+        <img src="https://zenodo.org/badge/XXXXXX.svg" alt="DOI"></a>
+    -->
+</p>
+
+This package provides decorators for memory utilization maximization with
+PyTorch and CUDA by starting with a maximum parameter size and applying
+successive halving until no more out-of-memory exception occurs.
+
+## 💪 Getting Started
+
+Assume you have a function for batched computation of nearest neighbors using
+brute-force distance calculation.
+
+```python
+import torch
+
+def knn(x, y, batch_size, k: int = 3):
+    return torch.cat(
+        [
+            torch.cdist(x[start : start + batch_size], y).topk(k=k, dim=1, largest=False).indices
+            for start in range(0, x.shape[0], batch_size)
+        ],
+        dim=0,
+    )
+```
+
+With `torch_max_mem` you can decorate this function to reduce the batch size
+until no more out-of-memory error occurs.
+
+```python
+import torch
+from torch_max_mem import maximize_memory_utilization
+
+
+@maximize_memory_utilization()
+def knn(x, y, batch_size, k: int = 3):
+    return torch.cat(
+        [
+            torch.cdist(x[start : start + batch_size], y).topk(k=k, dim=1, largest=False).indices
+            for start in range(0, x.shape[0], batch_size)
+        ],
+        dim=0,
+    )
+```
+
+In the code, you can now always pass the largest sensible batch size, e.g.,
+
+```python
+x = torch.rand(100, 100, device="cuda")
+y = torch.rand(200, 100, device="cuda")
+knn(x, y, batch_size=x.shape[0])
+```
+
+## 🚀 Installation
+
+The most recent release can be installed from
+[PyPI](https://pypi.org/project/torch_max_mem/) with uv:
+
+```console
+uv pip install torch_max_mem
+```
+
+or with pip:
+
+```console
+python3 -m pip install torch_max_mem
+```
+
+The most recent code and data can be installed directly from GitHub with uv:
+
+```console
+uv pip install git+https://github.com/mberr/torch-max-mem.git
+```
+
+or with pip:
+
+```console
+python3 -m pip install git+https://github.com/mberr/torch-max-mem.git
+```
+
+## 👐 Contributing
+
+Contributions, whether filing an issue, making a pull request, or forking, are
+appreciated. See
+[CONTRIBUTING.md](https://github.com/mberr/torch-max-mem/blob/master/.github/CONTRIBUTING.md)
+for more information on getting involved.
+
+## 👋 Attribution
+
+Parts of the logic have been developed with
+[Laurent Vermue](https://github.com/lvermue) for
+[PyKEEN](https://github.com/pykeen/pykeen).
+
+### ⚖️ License
+
+The code in this package is licensed under the MIT License.
+
+### 🍪 Cookiecutter
+
+This package was created with
+[@audreyfeldroy](https://github.com/audreyfeldroy)'s
+[cookiecutter](https://github.com/cookiecutter/cookiecutter) package using
+[@cthoyt](https://github.com/cthoyt)'s
+[cookiecutter-snekpack](https://github.com/cthoyt/cookiecutter-snekpack)
+template.
+
+## 🛠️ For Developers
+
+<details>
+  <summary>See developer instructions</summary>
+
+The final section of the README is for if you want to get involved by making a
+code contribution.
+
+### Development Installation
+
+To install in development mode, use the following:
+
+```console
+git clone git+https://github.com/mberr/torch-max-mem.git
+cd snekpack-demo
+uv pip install -e .
+```
+
+Alternatively, install using pip:
+
+```console
+python3 -m pip install -e .
+```
+
+### Updating Package Boilerplate
+
+This project uses `cruft` to keep boilerplate (i.e., configuration, contribution
+guidelines, documentation configuration) up-to-date with the upstream
+cookiecutter package. Install cruft with either `uv tool install cruft` or
+`python3 -m pip install cruft` then run:
+
+```console
+cruft update
+```
+
+More info on Cruft's update command is available
+[here](https://github.com/cruft/cruft?tab=readme-ov-file#updating-a-project).
+
+### 🥼 Testing
+
+After cloning the repository and installing `tox` with
+`uv tool install tox --with tox-uv` or `python3 -m pip install tox tox-uv`, the
+unit tests in the `tests/` folder can be run reproducibly with:
+
+```console
+tox -e py
+```
+
+Additionally, these tests are automatically re-run with each commit in a
+[GitHub Action](https://github.com/mberr/torch-max-mem/actions?query=workflow%3ATests).
+
+### 📖 Building the Documentation
+
+The documentation can be built locally using the following:
+
+```console
+git clone git+https://github.com/mberr/torch-max-mem.git
+cd snekpack-demo
+tox -e docs
+open docs/build/html/index.html
+```
+
+The documentation automatically installs the package as well as the `docs` extra
+specified in the [`pyproject.toml`](pyproject.toml). `sphinx` plugins like
+`texext` can be added there. Additionally, they need to be added to the
+`extensions` list in [`docs/source/conf.py`](docs/source/conf.py).
+
+The documentation can be deployed to [ReadTheDocs](https://readthedocs.io) using
+[this guide](https://docs.readthedocs.io/en/stable/intro/import-guide.html). The
+[`.readthedocs.yml`](.readthedocs.yml) YAML file contains all the configuration
+you'll need. You can also set up continuous integration on GitHub to check not
+only that Sphinx can build the documentation in an isolated environment (i.e.,
+with `tox -e docs-test`) but also that
+[ReadTheDocs can build it too](https://docs.readthedocs.io/en/stable/pull-requests.html).
+
+#### Configuring ReadTheDocs
+
+1. Log in to ReadTheDocs with your GitHub account to install the integration at
+   https://readthedocs.org/accounts/login/?next=/dashboard/
+2. Import your project by navigating to https://readthedocs.org/dashboard/import
+   then clicking the plus icon next to your repository
+3. You can rename the repository on the next screen using a more stylized name
+   (i.e., with spaces and capital letters)
+4. Click next, and you're good to go!
+
+### 📦 Making a Release
+
+#### Configuring Zenodo
+
+[Zenodo](https://zenodo.org) is a long-term archival system that assigns a DOI
+to each release of your package.
+
+1. Log in to Zenodo via GitHub with this link:
+   https://zenodo.org/oauth/login/github/?next=%2F. This brings you to a page
+   that lists all of your organizations and asks you to approve installing the
+   Zenodo app on GitHub. Click "grant" next to any organizations you want to
+   enable the integration for, then click the big green "approve" button. This
+   step only needs to be done once.
+2. Navigate to https://zenodo.org/account/settings/github/, which lists all of
+   your GitHub repositories (both in your username and any organizations you
+   enabled). Click the on/off toggle for any relevant repositories. When you
+   make a new repository, you'll have to come back to this
+
+After these steps, you're ready to go! After you make "release" on GitHub (steps
+for this are below), you can navigate to
+https://zenodo.org/account/settings/github/repository/mberr/torch-max-mem to see
+the DOI for the release and link to the Zenodo record for it.
+
+#### Registering with the Python Package Index (PyPI)
+
+You only have to do the following steps once.
+
+1. Register for an account on the
+   [Python Package Index (PyPI)](https://pypi.org/account/register)
+2. Navigate to https://pypi.org/manage/account and make sure you have verified
+   your email address. A verification email might not have been sent by default,
+   so you might have to click the "options" dropdown next to your address to get
+   to the "re-send verification email" button
+3. 2-Factor authentication is required for PyPI since the end of 2023 (see this
+   [blog post from PyPI](https://blog.pypi.org/posts/2023-05-25-securing-pypi-with-2fa/)).
+   This means you have to first issue account recovery codes, then set up
+   2-factor authentication
+4. Issue an API token from https://pypi.org/manage/account/token
+
+#### Configuring your machine's connection to PyPI
+
+You have to do the following steps once per machine.
+
+```console
+uv tool install keyring
+keyring set https://upload.pypi.org/legacy/ __token__
+keyring set https://test.pypi.org/legacy/ __token__
+```
+
+Note that this deprecates previous workflows using `.pypirc`.
+
+#### Uploading to PyPI
+
+After installing the package in development mode and installing `tox` with
+`uv tool install tox --with tox-uv` or `python3 -m pip install tox tox-uv`, run
+the following from the console:
+
+```console
+tox -e finish
+```
+
+This script does the following:
+
+1. Uses [bump-my-version](https://github.com/callowayproject/bump-my-version) to
+   switch the version number in the `pyproject.toml`, `CITATION.cff`,
+   `src/torch_max_mem/version.py`, and
+   [`docs/source/conf.py`](docs/source/conf.py) to not have the `-dev` suffix
+2. Packages the code in both a tar archive and a wheel using
+   [`uv build`](https://docs.astral.sh/uv/guides/publish/#building-your-package)
+3. Uploads to PyPI using
+   [`uv publish`](https://docs.astral.sh/uv/guides/publish/#publishing-your-package).
+4. Push to GitHub. You'll need to make a release going with the commit where the
+   version was bumped.
+5. Bump the version to the next patch. If you made big changes and want to bump
+   the version by minor, you can use `tox -e bumpversion -- minor` after.
+
+#### Releasing on GitHub
+
+1. Navigate to https://github.com/mberr/torch-max-mem/releases/new to draft a
+   new release
+2. Click the "Choose a Tag" dropdown and select the tag corresponding to the
+   release you just made
+3. Click the "Generate Release Notes" button to get a quick outline of recent
+   changes. Modify the title and description as you see fit
+4. Click the big green "Publish Release" button
+
+This will trigger Zenodo to assign a DOI to your release as well.
+
+</details>

torch_max_mem-0.1.4.dist-info/RECORD

@@ -0,0 +1,8 @@
+torch_max_mem/api.py,sha256=0576dc52db63f99c6bc928a4fb1bca2234391bb7d0276d8cbee8efa9c20e0c67,17786
+torch_max_mem/py.typed,sha256=e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855,0
+torch_max_mem/version.py,sha256=ee7ce9aee46d693e77190563fe86fd5a5544fda7bb7a61d95904d56b7c1c56b3,977
+torch_max_mem/__init__.py,sha256=1249af2acdc8d0ec62b94f722864cf85872502e9ccbc1cb97ad41b9025232e08,206
+torch_max_mem-0.1.4.dist-info/licenses/LICENSE,sha256=5163c26353a6aae35675c7d09095eab6c29149653ae83f942c510eab10e14c82,1071
+torch_max_mem-0.1.4.dist-info/WHEEL,sha256=a77f4251c0f6db962579c616a6b1d7cf825367be8db8f5d245046ce8ff723666,79
+torch_max_mem-0.1.4.dist-info/METADATA,sha256=29c64143887a54b549d60c8d1520ebbc674d411c1f4f9f123e0ee5338816a729,12714
+torch_max_mem-0.1.4.dist-info/RECORD,,

torch_max_mem-0.1.4.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.6.14
+Root-Is-Purelib: true
+Tag: py3-none-any

torch_max_mem-0.1.2.dist-info/METADATA

@@ -1,215 +0,0 @@
-Metadata-Version: 2.1
-Name: torch-max-mem
-Version: 0.1.2
-Summary: Maximize memory utilization with PyTorch.
-Home-page: https://github.com/mberr/torch-max-mem
-Download-URL: https://github.com/mberr/torch-max-mem/releases
-Author: Max Berrendorf
-Author-email: max.berrendorf@gmail.com
-Maintainer: Max Berrendorf
-Maintainer-email: max.berrendorf@gmail.com
-License: MIT
-Project-URL: Bug Tracker, https://github.com/mberr/torch-max-mem/issues
-Project-URL: Source Code, https://github.com/mberr/torch-max-mem
-Keywords: snekpack,cookiecutter,torch
-Classifier: Development Status :: 4 - Beta
-Classifier: Environment :: Console
-Classifier: Intended Audience :: Developers
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: OS Independent
-Classifier: Framework :: Pytest
-Classifier: Framework :: tox
-Classifier: Framework :: Sphinx
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3 :: Only
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: torch >=2.0
-Requires-Dist: typing-extensions
-Provides-Extra: docs
-Requires-Dist: sphinx <7 ; extra == 'docs'
-Requires-Dist: sphinx-rtd-theme ; extra == 'docs'
-Requires-Dist: sphinx-click ; extra == 'docs'
-Requires-Dist: sphinx-autodoc-typehints ; extra == 'docs'
-Requires-Dist: sphinx-automodapi ; extra == 'docs'
-Provides-Extra: formatting
-Requires-Dist: black ; extra == 'formatting'
-Requires-Dist: isort ; extra == 'formatting'
-Provides-Extra: tests
-Requires-Dist: numpy ; extra == 'tests'
-Requires-Dist: pytest ; extra == 'tests'
-Requires-Dist: coverage ; extra == 'tests'
-
-<!--
-<p align="center">
-  <img src="https://github.com/mberr/torch-max-mem/raw/main/docs/source/logo.png" height="150">
-</p>
--->
-
-<h1 align="center">
-  torch-max-mem
-</h1>
-
-<p align="center">
-    <a href="https://github.com/mberr/torch-max-mem/actions?query=workflow%3ATests">
-        <img alt="Tests" src="https://github.com/mberr/torch-max-mem/workflows/Tests/badge.svg" />
-    </a>
-    <a href="https://github.com/cthoyt/cookiecutter-python-package">
-        <img alt="Cookiecutter template from @cthoyt" src="https://img.shields.io/badge/Cookiecutter-snekpack-blue" /> 
-    </a>
-    <a href="https://pypi.org/project/torch_max_mem">
-        <img alt="PyPI" src="https://img.shields.io/pypi/v/torch_max_mem" />
-    </a>
-    <a href="https://pypi.org/project/torch_max_mem">
-        <img alt="PyPI - Python Version" src="https://img.shields.io/pypi/pyversions/torch_max_mem" />
-    </a>
-    <a href="https://github.com/mberr/torch-max-mem/blob/main/LICENSE">
-        <img alt="PyPI - License" src="https://img.shields.io/pypi/l/torch_max_mem" />
-    </a>
-    <a href='https://torch_max_mem.readthedocs.io/en/latest/?badge=latest'>
-        <img src='https://readthedocs.org/projects/torch_max_mem/badge/?version=latest' alt='Documentation Status' />
-    </a>
-    <a href='https://github.com/psf/black'>
-        <img src='https://img.shields.io/badge/code%20style-black-000000.svg' alt='Code style: black' />
-    </a>
-</p>
-
-This package provides decorators for memory utilization maximization with PyTorch and CUDA by starting with a maximum parameter size and applying successive halving until no more out-of-memory exception occurs.
-
-## 💪 Getting Started
-
-Assume you have a function for batched computation of nearest neighbors using brute-force distance calculation.
-
-```python
-import torch
-
-def knn(x, y, batch_size, k: int = 3):
-    return torch.cat(
-        [
-            torch.cdist(x[start : start + batch_size], y).topk(k=k, dim=1, largest=False).indices
-            for start in range(0, x.shape[0], batch_size)
-        ],
-        dim=0,
-    )
-```
-
-With `torch_max_mem` you can decorate this function to reduce the batch size until no more out-of-memory error occurs.
-
-```python
-import torch
-from torch_max_mem import maximize_memory_utilization
-
-
-@maximize_memory_utilization()
-def knn(x, y, batch_size, k: int = 3):
-    return torch.cat(
-        [
-            torch.cdist(x[start : start + batch_size], y).topk(k=k, dim=1, largest=False).indices
-            for start in range(0, x.shape[0], batch_size)
-        ],
-        dim=0,
-    )
-```
-
-In the code, you can now always pass the largest sensible batch size, e.g.,
-
-```python
-x = torch.rand(100, 100, device="cuda")
-y = torch.rand(200, 100, device="cuda")
-knn(x, y, batch_size=x.shape[0])
-```
-
-## 🚀 Installation
-
-The most recent release can be installed from
-[PyPI](https://pypi.org/project/torch_max_mem/) with:
-
-```bash
-$ pip install torch_max_mem
-```
-
-The most recent code and data can be installed directly from GitHub with:
-
-```bash
-$ pip install git+https://github.com/mberr/torch-max-mem.git
-```
-
-To install in development mode, use the following:
-
-```bash
-$ git clone git+https://github.com/mberr/torch-max-mem.git
-$ cd torch-max-mem
-$ pip install -e .
-```
-
-## 👐 Contributing
-
-Contributions, whether filing an issue, making a pull request, or forking, are appreciated. See
-[CONTRIBUTING.md](https://github.com/mberr/torch-max-mem/blob/master/CONTRIBUTING.md) for more information on getting involved.
-
-## 👋 Attribution
-
-Parts of the logic have been developed with [Laurent Vermue](https://github.com/lvermue) for [PyKEEN](https://github.com/pykeen/pykeen).
-
-
-### ⚖️ License
-
-The code in this package is licensed under the MIT License.
-
-### 🍪 Cookiecutter
-
-This package was created with [@audreyfeldroy](https://github.com/audreyfeldroy)'s
-[cookiecutter](https://github.com/cookiecutter/cookiecutter) package using [@cthoyt](https://github.com/cthoyt)'s
-[cookiecutter-snekpack](https://github.com/cthoyt/cookiecutter-snekpack) template.
-
-## 🛠️ For Developers
-
-<details>
-  <summary>See developer instrutions</summary>
-
-  
-The final section of the README is for if you want to get involved by making a code contribution.
-
-### 🥼 Testing
-
-After cloning the repository and installing `tox` with `pip install tox`, the unit tests in the `tests/` folder can be
-run reproducibly with:
-
-```shell
-$ tox
-```
-
-Additionally, these tests are automatically re-run with each commit in a [GitHub Action](https://github.com/mberr/torch-max-mem/actions?query=workflow%3ATests).
-
-### 📖 Building the Documentation
-
-```shell
-$ tox -e docs
-``` 
-
-### 📦 Making a Release
-
-After installing the package in development mode and installing
-`tox` with `pip install tox`, the commands for making a new release are contained within the `finish` environment
-in `tox.ini`. Run the following from the shell:
-
-```shell
-$ tox -e finish
-```
-
-This script does the following:
-
-1. Uses [Bump2Version](https://github.com/c4urself/bump2version) to switch the version number in the `setup.cfg` and
-   `src/torch_max_mem/version.py` to not have the `-dev` suffix
-2. Packages the code in both a tar archive and a wheel
-3. Uploads to PyPI using `twine`. Be sure to have a `.pypirc` file configured to avoid the need for manual input at this
-   step
-4. Push to GitHub. You'll need to make a release going with the commit where the version was bumped.
-5. Bump the version to the next patch. If you made big changes and want to bump the version by minor, you can
-   use `tox -e bumpversion minor` after.
-</details>

torch_max_mem-0.1.2.dist-info/RECORD

@@ -1,8 +0,0 @@
-torch_max_mem/__init__.py,sha256=7XoGfOMupwSZ3HWFXhwQB7ysIhPZeERbVRfzCDxaADw,230
-torch_max_mem/api.py,sha256=bTaDhXQIO8bv7xzL9TU-0IqqYeCRk66C5Ms70ZgcnFI,18228
-torch_max_mem/version.py,sha256=Bbj1Rx8IDqBJkeF-OjEagjvqA8K8k9XaQPSf5XQmVYU,1035
-torch_max_mem-0.1.2.dist-info/LICENSE,sha256=UWPCY1OmquNWdcfQkJXqtsKRSWU66D-ULFEOqxDhTII,1071
-torch_max_mem-0.1.2.dist-info/METADATA,sha256=ENMihj_gxuWWm7FbY5x_nc5NGlOpBOGlQQVqv3Y-Ew4,7401
-torch_max_mem-0.1.2.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-torch_max_mem-0.1.2.dist-info/top_level.txt,sha256=ztqqyZB7neLi8zWiWdaaKytNuRkZ7SzO7YDUQ4ZJR3U,14
-torch_max_mem-0.1.2.dist-info/RECORD,,

torch_max_mem-0.1.2.dist-info/WHEEL

@@ -1,5 +0,0 @@
-Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.2)
-Root-Is-Purelib: true
-Tag: py3-none-any
-

torch_max_mem-0.1.2.dist-info/top_level.txt

@@ -1 +0,0 @@
-torch_max_mem