Skip to main content

Functions to manipulate batches of PyTorch tensors

Project description

batchtensor

CI Nightly Tests Nightly Package Tests Codecov
Documentation Documentation
Code style: black Doc style: google Ruff Doc style: google
PYPI version Python BSD-3-Clause
Downloads Monthly downloads

Overview

batchtensor is lightweight library built on top of PyTorch to manipulate nested data structure with PyTorch tensors. This library provides functions for tensors where the first dimension is the batch dimension. It also provides functions for tensors representing a batch of sequences where the first dimension is the batch dimension and the second dimension is the sequence dimension.

Motivation

Let's imagine you have a batch which is represented by a dictionary with three tensors, and you want to take the first 2 items. batchtensor provides the function slice_along_batch that allows to slide all the tensors:

>>> import torch
>>> from batchtensor.nested import slice_along_batch
>>> batch = {
...     "a": torch.tensor([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
...     "c": torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]),
... }
>>> slice_along_batch(batch, stop=2)
{'a': tensor([[2, 6], [0, 3]]), 'b': tensor([4, 3]), 'c': tensor([1., 2.])}

Similarly, it is possible to split a batch in multiple batches by using the function split_along_batch:

>>> import torch
>>> from batchtensor.nested import split_along_batch
>>> batch = {
...     "a": torch.tensor([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
...     "c": torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]),
... }
>>> split_along_batch(batch, split_size_or_sections=2)
({'a': tensor([[2, 6], [0, 3]]), 'b': tensor([4, 3]), 'c': tensor([1., 2.])},
 {'a': tensor([[4, 9], [8, 1]]), 'b': tensor([2, 1]), 'c': tensor([3., 4.])},
 {'a': tensor([[5, 7]]), 'b': tensor([0]), 'c': tensor([5.])})

Please check the documentation to see all the implemented functions.

Documentation

  • latest (stable): documentation from the latest stable release.
  • main (unstable): documentation associated to the main branch of the repo. This documentation may contain a lot of work-in-progress/outdated/missing parts.

Installation

We highly recommend installing a virtual environment. batchtensor can be installed from pip using the following command:

pip install batchtensor

To make the package as slim as possible, only the minimal packages required to use batchtensor are installed. To include all the dependencies, you can use the following command:

pip install batchtensor[all]

Please check the get started page to see how to install only some specific dependencies or other alternatives to install the library. The following is the corresponding batchtensor versions and tested dependencies.

batchtensor coola numpy* torch python
main >=0.9.1,<1.0 >=1.24,<2.0 >=2.4,<3.0 >=3.10,<3.15
0.1.1 >=0.9.1,<1.0 >=1.24,<2.0 >=2.4,<3.0 >=3.10,<3.15
0.1.0 >=0.8.6,<1.0 >=1.21,<2.0 >=2.4,<3.0 >=3.9,<3.14
0.0.5 >=0.8.6,<1.0 >=1.21,<2.0 >=1.11,<3.0 >=3.9,<3.14
0.0.4 >=0.1,<1.0 >=1.21,<2.0 >=1.11,<3.0 >=3.9,<3.13
0.0.3 >=0.1,<1.0 >=1.21,<2.0 >=1.11,<3.0 >=3.9,<3.13
0.0.2 >=0.1,<1.0 >=1.21,<2.0 >=1.11,<3.0 >=3.9,<3.13
0.0.1 >=0.1,<0.4 >=1.21,<2.0 >=1.11,<3.0 >=3.9,<3.13

* indicates an optional dependency

Contributing

Please check the instructions in CONTRIBUTING.md.

Suggestions and Communication

Everyone is welcome to contribute to the community. If you have any questions or suggestions, you can submit Github Issues. We will reply to you as soon as possible. Thank you very much.

API stability

:warning: While batchtensor is in development stage, no API is guaranteed to be stable from one release to the next. In fact, it is very likely that the API will change multiple times before a stable 1.0.0 release. In practice, this means that upgrading batchtensor to a new version will possibly break any code that was using the old version of batchtensor.

License

batchtensor is licensed under BSD 3-Clause "New" or "Revised" license available in LICENSE file.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

batchtensor-0.1.2a0.tar.gz (132.6 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

batchtensor-0.1.2a0-py3-none-any.whl (38.4 kB view details)

Uploaded Python 3

File details

Details for the file batchtensor-0.1.2a0.tar.gz.

File metadata

  • Download URL: batchtensor-0.1.2a0.tar.gz
  • Upload date:
  • Size: 132.6 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.9.9 {"installer":{"name":"uv","version":"0.9.9"},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for batchtensor-0.1.2a0.tar.gz
Algorithm Hash digest
SHA256 44d73707b176a2021e4412cbb7a479ecef99a12a1123dee006f433d8e8f28f33
MD5 3113eeedc0bd72e3b4d240eb131d6829
BLAKE2b-256 97daab7b569e77e24a2f89d02987777d161d6ec6436b984e2539c234abeb771f

See more details on using hashes here.

File details

Details for the file batchtensor-0.1.2a0-py3-none-any.whl.

File metadata

  • Download URL: batchtensor-0.1.2a0-py3-none-any.whl
  • Upload date:
  • Size: 38.4 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.9.9 {"installer":{"name":"uv","version":"0.9.9"},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for batchtensor-0.1.2a0-py3-none-any.whl
Algorithm Hash digest
SHA256 d299935f1f78d5a7d83232178ee955724454557fb9df318eb61014a16cde855c
MD5 7135432c0b57f635c5666c19a26993fa
BLAKE2b-256 cbbe967b4b2ddf38522c7d1f9f0704bb66b53782271a62d78af2ce739b2eef14

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page