Skip to content

Add Python APIs for PTD tensor dictionaries - #21236

Open
Kilbex wants to merge 2 commits into
pytorch:mainfrom
Kilbex:python-ptd-tensor-dicts
Open

Add Python APIs for PTD tensor dictionaries#21236
Kilbex wants to merge 2 commits into
pytorch:mainfrom
Kilbex:python-ptd-tensor-dicts

Conversation

@Kilbex

@Kilbex Kilbex commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds public Python helpers for saving and loading PTD tensor dictionaries. The helpers validate PTD-compatible CPU tensors, preserve supported strides, normalize other layouts, and reject malformed tensor entries so callers no longer need to construct DataPayload objects directly.

Test plan

PYTHONPATH=.. python -m pytest extension/flat_tensor/test/test_serialize.py -q

lintrunner on the changed files

Authored with Claude.

@pytorch-bot

pytorch-bot Bot commented Jul 23, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21236

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 3e30486 with merge base 430b73d (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 23, 2026
@Kilbex

Kilbex commented Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

@pytorchbot label "release notes: api"

@pytorch-bot pytorch-bot Bot added the release notes: api Changes to public facing apis (any interfaces, pybinded runtime methods, etc.) label Jul 23, 2026
@nil-is-all nil-is-all added the module: extension Issues related to code under extension/ label Jul 23, 2026
@nil-is-all

Copy link
Copy Markdown
Contributor

Thanks for the PR, @Kilbex. Waiting on CI to finish running; will keep you posted

@shoumikhin shoumikhin left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, this is clean and the tricky tensor-memory cases are handled well (stride-0 expands, slices of a larger storage, empty tensors, writable frombuffer). Tests are solid. One thing I want to confirm before stamping, plus two smaller notes.

1. alignment=1 on every saved entry. save_ptd writes every entry with alignment=1. These .ptd files are also loaded by the C++ runtime, which normally relies on segment alignment for mmap / zero-copy loads (the usual serialize path uses the FlatTensorConfig alignment, not 1). Can you confirm a file written by save_ptd actually loads correctly and efficiently in the runtime, not just in a Python round-trip? If these are meant to be Python-only, please say so in the docstring/README. Otherwise I think this should use the same default alignment as the normal path.

2. dtype table import robustness. _PTD_TO_TORCH_DTYPE is built at import time and references several newer dtypes (torch.uint16/uint32/uint64, torch.bits16, the float8_*fnuz variants, quint4x2/2x4). On a torch build that lacks any of them, importing this module fails with AttributeError. Consider guarding with getattr(torch, name, None) and skipping missing ones, so the import stays robust across torch versions.

3. Document the lossy layout normalization. Worth making explicit that a non-contiguous / non-channels-last tensor round-trips with the same values and shape but contiguous strides (the layout is normalized, not preserved). The README hints at it; a one-line note on save_ptd so callers do not expect stride-identical round-trips would help.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: extension Issues related to code under extension/ release notes: api Changes to public facing apis (any interfaces, pybinded runtime methods, etc.)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants