Add Python APIs for PTD tensor dictionaries - #21236
Conversation
🔗 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 FailuresAs of commit 3e30486 with merge base 430b73d ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot label "release notes: api" |
|
Thanks for the PR, @Kilbex. Waiting on CI to finish running; will keep you posted |
shoumikhin
left a comment
There was a problem hiding this comment.
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.
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
DataPayloadobjects directly.Test plan
PYTHONPATH=.. python -m pytest extension/flat_tensor/test/test_serialize.py -qlintrunneron the changed filesAuthored with Claude.