🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Run pull request jobs on OSDC runners in shadow mode

❌ 1 Unclassified Failure

As of commit df04807 with merge base f513851 ():

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

@claude review these changes

Claude finished @jansel's task in 1m 32s —— View job

Review of PR #183756: Stabilize jacobian helper FX graphs

• Gather context and read the diff
• Review torch/_functorch/eager_transforms.py changes
• Review test/dynamo/test_higher_order_ops.py changes
• Post review feedback

Overall: LGTM — this is a clean, correct change that makes the jacobian FX graphs more stable for Dynamo.

Summary of changes

The PR does two things:

1. Replaces CPU tensor bookkeeping with plain Python in _chunked_standard_basis_for_ (lines 900-904): The original code used torch.tensor(tensor_numels).cumsum(dim=0)[:-1].neg().unbind() to compute diagonal start indices, which produced tensor operations that appeared in Dynamo-traced graphs. The new code computes the same values via a simple Python loop, so these operations no longer pollute the FX graph.

2. Uses torch.split (functional form) instead of Tensor.split (method form) in jacrev (line 781) and jacfwd (line 1409): This ensures the traced graph always spells the op as torch.functional.split(...) rather than .split(...), eliminating method vs. function spelling instability across traces.

Correctness verification

The diagonal offset computation is semantically equivalent:

• Original: (0, *torch.tensor([a, b, c]).cumsum(0)[:-1].neg().unbind()) → (0, -a, -(a+b))
• New: iterative subtraction → [0, -a, -(a+b)]

Both produce the same sequence of offsets.

Minor observations

• The list() wrapping in torch.split(result, list(flat_output_numels), dim=0) is needed because flat_output_numels is a tuple (from tuple(out.numel() for out in flat_output) at line 662), and torch.split with a sequence argument expects a list for the split sections. This is fine.

• In test_jacrev_two_tensors_argnums (around line 4107-4162), the assertion was restructured to use empty_line_normalizer on both sides to handle cross-platform empty line differences — a pragmatic fix for test stability.

• The test changes are purely mechanical: removing the now-absent tensor/cumsum/getitem/neg/unbind nodes and updating the split spelling from .split((12,), dim=...) to torch.functional.split(..., [12], dim=...).

No issues found. The change is minimal, well-targeted, and directly addresses the graph instability reported in #180264.