[BugFix] Implement pytree flatten_with_keys for lazy stacks

[YeonwooSung]

Problem

_lazy_td_flatten_with_keys was registered as the pytree flatten_with_keys_fn for LazyStackedTensorDict, but it only raised NotImplementedError. tree_flatten_with_path (and tree_map_with_path) therefore failed on a lazy stack.

torch.export compares the spec of tree_flatten_with_path on the call inputs with the spec of tree_flatten on the example inputs. Dense tensordicts already keep those specs aligned by returning the same context from _td_flatten_with_keys as from _tensordict_flatten. A lazy stack did not.

Fix

_lazy_td_flatten_with_keys now calls _lazy_tensordict_flatten and returns that same context, with each member wrapped in SequenceKey(i) in member order. Unflatten is unchanged.

Tests

TestPyTree.test_pytree_lazy_flatten_with_keys checks that:

• tree_flatten_with_path does not raise on a lazy stack
• its spec and context match tree_flatten
• keys are SequenceKey in member order, including a nested member and an exclusive key
• tree_unflatten rebuilds a lazy stack with the same stack_dim, values, names, and batch size
• a tensorclass member roundtrips as a tensorclass

Before the fix the new test failed with NotImplementedError. After the fix:

pytest test/nn/test_functorch.py::TestPyTree::test_pytree_lazy_flatten_with_keys \
       test/nn/test_functorch.py::TestPyTree::test_pytree_lazy \
       test/nn/test_functorch.py::TestPyTree::test_spec_batch_size -q --tb=short
3 passed

TestPyTree (15 tests) also passed. Ruff check and format are clean on the touched files.