Adding support for calling `value_and_grad` on non-scalar inputs

[paul0403]

Before submitting

Please complete the following checklist when submitting a PR:

• [x] All new functions and code must be clearly commented and documented.

• [x] Ensure that code is properly formatted by running make format. The latest version of black and clang-format-14 are used in CI/CD to check formatting.

• [x] All new features must include a unit test. Integration and frontend tests should be added to frontend/test, Quantum dialect and MLIR tests should be added to mlir/test, and Runtime tests should be added to runtime/tests.

When all the above are checked, delete everything above the dashed line and fill in the pull request template.

---

Context: Bug fix: currently catalyst.value_and_grad only works on scalar inputs, from the work done in #804. This PR extends to arbitrary pytree shapes.

Description of the Change:

• In frontend, record the input pytree shape for the gradient calls in differentiation.py amd jax_primitives.py for value_and_grad
• In both frontend and mlir, remove the assertion checking the total number of outputs is even, since the differentiated function can have any pytree shape as input and hence can have any number of inputs
• In mlir, when verifying a ValueAndGradOp, check the gradient result types against the callee input argument types from the right instead of from the left, since if the function have constant values, the cloned callee will take in these constant values as arguments at the start of the argument list
• Added frontend tests for value_and_grad

Benefits: catalyst.value_and_grad now can take in functions whose numerical arguments are not scalars

Possible Drawbacks:

Related GitHub Issues: closes #841 [sc-66764]

[josh146]

Thanks @paul0403! This PR should be made to the v0.7.0-rc branch, since it's a bugfix to a release feature :) cc @rauletorresc

[rauletorresc]

Thanks @paul0403! This PR should be made to the v0.7.0-rc branch, since it's a bugfix to a release feature :) cc @rauletorresc

Thanks for the hint! I have rebased it with respect to v0.7.0-rc

[josh146]

Also dow we plan to accept pytree values as return for example {"result": x}? It does not seem very useful as we are limited to one value returned.

For return I don't mind either way, but definitely important to support as input (as I've seen JAX optimizers that require this)

[josh146]

Although I think @jax.value_and_grad works with returned pytrees, so I think it's good to match it's behaviour.

[josh146]

@paul0403 don't forget to add a link to this PR to the changelog! You don't need to create a new changelog entry, simply add this PR to the existing value_and_grad changelog entry.

[paul0403]

Nice job! Just a question about the assertions. Also dow we plan to accept pytree values as return for example {"result": x}? It does not seem very useful as we are limited to one value returned.

Currently catalyst only supports the differentiated function having a single float as return; this is not a catalyst.value_and_grad issue, as it is the behavior of catalyst.grad as well.