Rebuild the full argument list when partitioning activations

[vineethsaivs]

With partition_activations on, a checkpointed function that takes any non-tensor argument gets a corrupted argument list on the recompute.

Cause: get_partitioned_activations_for_backward saves a (value, size) pair for every argument, so tensor_flags and non_tensor_objects both hold two entries per argument. merge_tensors collapsed the flags by skipping the one after a True, which drops nothing for a non-tensor argument, whose value and its None size are both non-tensors. It never collapsed the non-tensor values at all.

Driving the real save and restore helpers on CPU:

forward was called with	recompute received
(tensor, tensor)	(tensor, tensor)
(tensor, 2)	(tensor, 2, None)
(tensor, None)	(tensor, None, None)
(tensor, None, 2.5)	(tensor, None, None, 2.5, None)
(tensor, None, True, tensor)	(tensor, None, None, True, None, tensor)

So a stray None lands after every non-tensor argument and shifts everything after it. TestCheckpointNonTensor already covers these exact argument shapes, but never with partition_activations.

Fix: every argument owns an even index in both lists, so drop the odd ones.

real_tensor_flags = tensor_flags[::2]
non_tensor_objects = non_tensor_objects[::2]

All-tensor arguments are unaffected, which is why this went unnoticed.

Test: two tests in tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py drive the real get_partitioned_activations_for_backward, extract_tensors, gather_partitioned_activations and merge_tensors, so they run on CPU without an accelerator.

pytest tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py -k partitioned
  on master:  6 failed  (assert 3 == 2)
  with this:  6 passed

The end-to-end path cannot run here: _test_activation_checkpoint skips on the CPU accelerator, and CheckpointFunction.forward needs a real Stream. Flagging that rather than implying otherwise. yapf and flake8 --config=.flake8 are clean on both files.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f0f4776096

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add a device-backed integration regression

When partitioned activation checkpointing is used in either public checkpoint path, these tests only invoke internal serialization helpers and deliberately bypass CheckpointFunction.forward/backward and the non-reentrant unpack hook. They therefore do not establish that the changed training-loop contract works through the actual checkpoint execution path. Add and run a minimal device-backed training-loop regression using checkpoint(..., partition_activations=True) with non-tensor arguments, comparing outputs and gradients to an uncheckpointed run.

AGENTS.md reference: AGENTS.md:L35-L35

Useful? React with 👍 / 👎.