[Capture] Allow higher order primitives to accept dynamically shaped arrays

[albi3ro]

Context:

By turning on the experimental jax_dynamic_shapes mode, you can capture and compile jaxpr for a series of different shapes at the same time. While this expermental feature has issues and isn't fully supported by jax yet, it is used by catalyst. To continue to support all of catalyst's features, we need to be able to capture and work with dynamic shapes as well.

Description of the Change:

• Adds a qml.capture.determine_abstracted_axes function to determine the required abstracted_axes and the corresponding abstract shapes.
• Use the determine_abstracted_axes function in all of our higher order primitives other than grad and jacobian, as grad and jacobian may prove more complicated.
• Add a document explaining abstract shapes and how we can work with them.

Benefits:

Our higher order primitives can accept inputs with abstract shapes.

Possible Drawbacks:

This jax mode is still experimental.

Related GitHub Issues:

[sc-81471]

[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 ✅

Project coverage is 99.54%. Comparing base (afec979) to head (35b0091).
Report is 2 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6786   +/-   ##
=======================================
  Coverage   99.54%   99.54%           
=======================================
  Files         477      478    +1     
  Lines       45246    45296   +50     
=======================================
+ Hits        45042    45092   +50     
  Misses        204      204           

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

[lillian542]

Left a few minor comments, but I have questions about determine_abstract_axes 😅

I need some help understanding what is going on, it would probably be easiest just to chat.

[mudit2812]

Looks good. It seems clear to me that "dynamic shapes" at this point refers to dynamic axis sizes, but the number of dimensions cannot be dynamic? Let me know if I'm misinterpreting. If not, probably worth noting that in the markdown file.

[albi3ro]

Looks good. It seems clear to me that "dynamic shapes" at this point refers to dynamic axis sizes, but the number of dimensions cannot be dynamic? Let me know if I'm misinterpreting. If not, probably worth noting that in the markdown file.

@mudit2812 Yes, the number of dimensions cannot be dynamic. Only the size of the dimensions.

[mudit2812]

Looks great! Just a few final clean up related comments, otherwise I'm happy to approve.