[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.60%. Comparing base (807bc4c) to head (47bdf11).
Report is 2 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6786   +/-   ##
=======================================
  Coverage   99.60%   99.60%           
=======================================
  Files         476      477    +1     
  Lines       45237    45282   +45     
=======================================
+ Hits        45060    45105   +45     
  Misses        177      177           

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

[mudit2812]

Don't forget to add your name to the changelog 😄

[albi3ro]

added. thanks for the reminder.

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

[lillian542]

Suggested change

	Dynamic shapes are experimental feature of jax with limited support and feature coverage.
	Dynamic shapes are an experimental feature of jax with limited support and feature coverage.

[lillian542]

This really confused me for a moment because I read it as "without the above setup", and I thought it implied we need to run jax.config.update("jax_dynamic_shapes", False) in order for the dynamic arrays to work.

Suggested change

	```python
	jax.config.update("jax_dynamic_shapes", False)
	```
	Without this setup, we can't create arrays whose size depends on an abstract value.
	Without the `"jax_dynamic_shapes"` feature, we can't create arrays whose size depends on an abstract value.
	```python
	jax.config.update("jax_dynamic_shapes", False)

[lillian542]

😢

[lillian542]

I'm a little lost on the alphabet part

[lillian542]

Suggested change

	We can now take these learnings a make custom higher order primitive that supports dynamically shaped inputs:
	We can now take these learnings to make a custom higher order primitive that supports dynamically shaped inputs:

[lillian542]

🤣

[lillian542]

Below we have *abstract_shapes, *flat_args in some of the primitive bind functions, and here it's reversed; it would be more consistent to have a set order.

[lillian542]

I'm having a hard time reviewing this function because I'm not confident that I understand what the expected behaviour is for various inputs. Maybe we could have a quick call/meeting and you could walk me through it?