Skip to main content

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.41

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.41
File Size Uploaded Provenance
torch_einops_utils-0.1.41.tar.gz 31.0 kB Without provenance Details

Built distribution (wheel)

Table of built distributions (wheels) for torch-einops-utils 0.1.41
File Interpreter ABI Platform
torch_einops_utils-0.1.41-py3-none-any.whl Python 3 none any Details

Total release size: 51.2 kB

Release files / torch_einops_utils-0.1.41.tar.gz

Download URL torch_einops_utils-0.1.41.tar.gz
Size 31.0 kB
Tags Source
SHA-256 checksum
How to use checksums
a7bdae99906cac1e988e5909302225e4a9735d3e386c9d856b264d8fe20cf41c
BLAKE2b-256 checksum
How to use checksums
615bb2e1ad51eb901d203a697e9bda43a3ef630bc676c2e67fc938c67693dd81
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via uv/0.8.17

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

Download URL torch_einops_utils-0.1.41-py3-none-any.whl
Size 20.1 kB
Tags Python 3
SHA-256 checksum
How to use checksums
2df38e0d21940868af745d1f3898a3bfabf1fd6626ff838e96cf9acd7fb5072c
BLAKE2b-256 checksum
How to use checksums
3c306436c26428f5107d4c26d9a560a9f3fae92c4b3a58fa925fbb8b507b1615
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via uv/0.8.17

Security for torch-einops-utils 0.1.41

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

No provenance found

None of the files in this release were uploaded with provenance metadata.

A note on package security

PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.

To learn more about protecting yourself from malicious actors, visit our package security documentation.

Release history Release notifications | RSS feed

This release

0.1.41 This release

2 release files No provenance

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page