Add Python save_ptd/load_ptd for tensor dicts

[akashjainn]

Fixes #8542.

ExecuTorch has save_ptd in C++ (extension/flat_tensor/serialize/serialize.h, used by the CIFAR and XOR training examples) but nothing equivalent in Python - grep -rn "save_ptd|load_ptd" --include=*.py comes back empty. The only Python entry point is FlatTensorSerializer, which takes a DataPayload rather than a plain dict[str, torch.Tensor].

Per @JacobSzwejbka's comment on the issue, the ask is round-tripping dict[str, tensor], which #11779 made straightforward.

Added save_ptd and load_ptd in extension/flat_tensor/serialize/serialize.py, mirroring the C++ argument order and taking either a path or a binary file-like object, which covers both C++ overloads in one function. Pure Python on top of the existing FlatTensorSerializer, so no new pybind11 surface and no second implementation to drift.

save_ptd(path, tensor_map, tensor_alignment=16)   # default matches the C++ call sites
load_ptd(path) -> Dict[str, torch.Tensor]

A few things that needed care:

bfloat16 and non-contiguous layouts - tensor bytes go through the existing _tensor_to_bytes helper rather than a fresh .numpy().tobytes(), since bfloat16 has no numpy dtype.

channels_last - dim_order is a real permutation, so the bytes on disk aren't in logical order. load_ptd reshapes to the physical sizes implied by dim_order and then inverts that permutation. Without it the tensor comes back transposed. There's a test asserting both value equality and that the memory format survives.

Tied weights - a tensor under several keys, like an embedding reused as the output projection, is written once with both keys pointing at one buffer.

Empty tensors are built directly, since torch.frombuffer rejects a zero-length buffer. Named data entries with no tensor layout aren't tensors and get skipped rather than mis-decoded.

Ten new tests: dtype round-trips including bfloat16, scalar and empty tensors, channels_last, tied-tensor dedup, the file-object overload, alignment of segment offsets, and the two error paths. All pass, along with the existing TestSerialize cases in the same file.

ruff isn't used here; ufmt (black 24.4.2 + usort) and flake8 6.1.0 are clean on the changed files.

Note that extension/flat_tensor/test/test_serialize.py is in the --ignore list in pytest.ini (pre-existing, for missing flatc), so these run through the internal python_unittest target rather than OSS pytest. I verified them locally against a flatc-provisioned environment.

Two things I'd like your call on. I mirrored the C++ argument order with path first, since the issue framed this as the Python form of the C++ API - the Python-native alternative would be save_ptd(tensor_map, path, ...) matching torch.save(obj, f). Happy to flip it. And I deliberately left extension/flat_tensor/init.py empty; populating it would give a tidier import, but no python_library target currently covers that file and it would make importing the subpackage pull in torch. Glad to add it if you'd prefer the shorter import.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 15 Awaiting Approval

As of commit a9864f7 with merge base c8d5189 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: akashjainn / name: akashjainn (a9864f7)

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[akashjainn]

/easycla