perf(transforms): Use stop_gradient in custom_function VJP

[Rudra-Tiwari-codes]

Proposed changes

This PR optimizes the fallback VJP logic within custom_function in mlx/transforms.cpp. Previously, the VJP computation tracked gradients for all inputs, even those not specified in argnums. This PR introduces stop_gradient for inputs that are not part of the differentiation request. This change:

Prevents wasteful computation of gradients for constant or irrelevant inputs.
Ensures that argnums is fully respected, aligning with the expected behavior of efficient gradient tracking.

Checklist

Put an x in the boxes that apply.

• [ x] I have read the CONTRIBUTING document
• [ x] I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• [ x] I have added tests that prove my fix is effective or that my feature works
• [ x] I have updated the necessary documentation (if needed)

[fransafu]

Warning: This author is forking multiple ML projects such as google-deepmind/alphafold, ml-explore/mlx, openai/CLIP, pytorch/pytorch, tensorflow/tensorflow, anthropics/claude-code, vllm-project/vllm, and others, adding minimal "contributions" (often for tests or miscellaneous changes) without proper validation. A review of their commits shows mostly local implementations of TODOs copied from existing projects, with little to no substantive review or testing.

So far, this author has forked 41 repositories following the same pattern. Be careful when accepting this PR. It’s also concerning how this author is able to submit PRs across four repositories in the same day, each requiring large context, which strongly suggests a highly automated workflow.