BugFix: adjoint metric tensor with jax

[astralcai]

Context:
The adjoint_metric_tensor transform does not work with jax variables because jax variables are not considered trainable parameters until they become tracers.

Description of the Change:

1. Add a custom expand transform to adjoint_metric_tensor that expands trainable parameters based on argnums.
2. Add an optional use_argnum_in_expand argument to transform programs to determine whether or not to perform jax_argnums_to_tape_trainable on the parameters.

Benefits:
BugFix

Possible Drawbacks:
Adding yet another keyword argument to transform may make code look messy.

Related GitHub Issues:
#5197

Related Shortcut Stories:
[sc-56734]

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

❗ No coverage uploaded for pull request base (v0.35.0-rc0@c1f3997). Click here to learn what that means.

Additional details and impacted files

@@              Coverage Diff               @@
##             v0.35.0-rc0    #5271   +/-   ##
==============================================
  Coverage               ?   99.65%           
==============================================
  Files                  ?      399           
  Lines                  ?    36617           
  Branches               ?        0           
==============================================
  Hits                   ?    36489           
  Misses                 ?      128           
  Partials               ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[rmoyard]

It looks good! I don't think it is necessary to argnums as arguments though

[rmoyard]

Thanks @astralcai , it looks good to me 💯

[timmysilv]

looks good!