Fix KeyError for args_tensor_mask removed in PyTorch nightly

[IvanYashchuk]

A recent PyTorch PR #166788 changed the autograd_function_apply higher-order op API by removing the args_tensor_mask keyword argument.

Currently, Thunder reads this key in the fwd_kwargs dict leading to a KeyError:

lightning-thunder/thunder/core/jit_ext.py

Line 943 in fb989d4

args_tensor_mask = unwrap(fwd_kwargs["args_tensor_mask"])

This PR adds a compatibility layer for both older PyTorch releases (which still provide args_tensor_mask) and newer PyTorch nightly versions (which removed it). When the key is not present, we infer the tensor mask by checking which arguments are TensorProxy instances.

Let's see in CI if further changes needed anywhere else.

Fixes #2803.

[Copilot]

Pull request overview

This PR adds a compatibility layer to handle the removal of the args_tensor_mask keyword argument from PyTorch's autograd_function_apply higher-order op API (removed in PyTorch PR #166788). The fix ensures Thunder works with both older PyTorch versions (that provide args_tensor_mask) and newer nightly versions (that removed it).

Key changes:

• Added conditional logic to check if args_tensor_mask exists in fwd_kwargs before accessing it
• When not present, infers the tensor mask by checking which arguments are TensorProxy instances

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The comment states "When not present, we assume all fwd_args are tensors" but the implementation actually checks each argument to determine if it's a TensorProxy. The comment should be updated to accurately reflect what the code does: "When not present, we infer the tensor mask by checking if each argument is a TensorProxy".

Suggested change

	# When not present, we assume all fwd_args are tensors.
	# When not present, we infer the tensor mask by checking if each argument is a TensorProxy.

[IvanYashchuk]

Closing in favor of #2802.