torch-einops-utils

Some utility functions to help myself (and perhaps others) go faster with ML/AI work

maybe_return

return just the primary output by default, extras only on request

from torch_einops_utils import maybe_return

@maybe_return('memories')
def forward(self, ids, return_memories = False):
    memories = build_memories(ids) if return_memories else None

    return backbone(ids), memories

model(ids)                             # logits
model(ids, return_memories = True)     # Output(logits, memories)
logits, memories = model(ids, return_memories = True)  # unpacks

any number of extras, individually or all at once

@maybe_return('memories', 'hiddens')
def forward(self, ids, return_memories = False, return_hiddens = False):
    ...
    return logits, memories, hiddens

model(ids, return_hiddens = True)  # Output(logits, hiddens)
model(ids, return_all = True)      # Output(logits, memories, hiddens)

primary may be a tuple of names, or a callable receiving the bound call arguments, for when the shape of the primary depends on the mode

def primary_fields(kwargs):
    if not kwargs.get('return_loss', False):
        return ('logits', 'memory') if kwargs.get('return_memory', False) else 'logits'

    return ('loss', 'loss_mask') if not kwargs.get('reduce_loss', False) else 'loss'

@maybe_return('pooled_repr', primary = primary_fields)
def forward(self, ids, return_loss = False, return_memory = False, reduce_loss = False, return_pooled_repr = False):
    ...
    if return_loss:
        return loss, loss_mask, pooled_repr

    return logits, memory, pooled_repr

model(ids)                            # logits
logits, memory = model(ids, return_memory = True)
loss, loss_mask = model(ids, return_loss = True)
loss = model(ids, return_loss = True, reduce_loss = True)
loss, pooled = model(ids, return_loss = True, reduce_loss = True, return_pooled_repr = True)

Metadata

Release files for torch-einops-utils 0.1.36

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for torch-einops-utils 0.1.36
File	Size	Uploaded
torch_einops_utils-0.1.36.tar.gz	28.3 kB	Oct 5, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for torch-einops-utils 0.1.36
File	Interpreter	ABI	Platform	Reset
torch_einops_utils-0.1.36-py3-none-any.whl 18.4 kB Oct 5, 2026	Python 3	none	any	Details

Total release size: 46.7 kB

Release files / torch_einops_utils-0.1.36.tar.gz

Download URL	torch_einops_utils-0.1.36.tar.gz
Size	28.3 kB
Tags	Source
SHA-256 checksum How to use checksums	`106ae495450e0a9ec3fb150a2f1d1d36628f1a5884abad2bf63fce240551ff14`
BLAKE2b-256 checksum How to use checksums	`7a3aaa3cf1eef4a7305f29df5961d1e573e8435b00453dadd131e8153bae81ff`
Upload date	Oct 5, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`

Release files / torch_einops_utils-0.1.36-py3-none-any.whl

Download URL	torch_einops_utils-0.1.36-py3-none-any.whl
Size	18.4 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`1806455c4c76318cf11e604fa9501187be32eebce8ffd6973e1c52336f34f4cb`
BLAKE2b-256 checksum How to use checksums	`2e5ddac32f7b925b4a7b1ba0a1f1dff392c96f77910a7f1ba203d185e749c52a`
Upload date	Oct 5, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`