[Frontend] Support capturing dynamically-shaped arrays from outer scopes in Python programs

[sergei-mironov]

Context: Catalyst supports Jax dynamically-shaped arrays. The current version has a notable limitation: body loop programs do not allow mixing captured dynamically-shaped arrays with the argument ones even if they are of the same dimension. This is because we effectively duplicated dimension variables for the loop-body arguments.

As an illustration, the loop body of the below program takes a as an argument array and x as a captured array. The experimental_preserve_dimensions flag has the default value of True

@qjit(abstracted_axes={1: 'n'})
def g(x, y):

  @catalyst.for_loop(0, 10, 1, experimental_preserve_dimensions=True)
  def loop(_, a):
    return a * x

  return jnp.sum(loop(y))

a = jnp.ones([1,3], dtype=float)
b = jnp.ones([1,3], dtype=float)
g(a, b)

Description of the Change:

This PR sets the following semantic of loops, depending on the value of the already-existing experimental_preserve_dimensions flag:

• True (the default): all dynamic dimension variables will be handled as Jaxpr constants. So
  • Mixing argument and captured arrays is possible.
  • No dimension modification is allowed inside loops.
• False would handle all the dimensions as Jaxpr implicit arguments.
  • Mixing argument and captured arrays is not possible.
  • Dimension modification within the loops are allowed.

[sc-60521]

Benefits:

Possible Drawbacks:

Related GitHub Issues:

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 97.97%. Comparing base (23e6d2b) to head (9b0277e). Report is 207 commits behind head on main.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #830 +/- ## ========================================== - Coverage 97.98% 97.97% -0.02% ========================================== Files 71 71 Lines 10429 10445 +16 Branches 936 947 +11 ========================================== + Hits 10219 10233 +14 - Misses 168 169 +1 - Partials 42 43 +1 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[josh146]

Hi @rauletorresc! Since this fixes a bug in the dynamic shape array in v0.7.0, do we want to cherry pick this into the v0.7.0-rc branch?

cc @dime10 and @erick-xanadu for thoughts.

[rauletorresc]

Makes sense. I'll proceed once the discussion is finished :)

[erick-xanadu]

@josh146 : what bug is fixed?

[josh146]

@erick-xanadu it is a bit blurry, but I view the current issue with dynamic shape support in v0.7.0 (that it doesn't support capturing dynamic shapes from outer scopes in python programs) as a bug.

So would be in favor of this being merged in as a bug fix to that behaviour in v0.7.0-rc.

Since this has already been merged into main, we will need to cherry pick this commit into v0.7.0-rc.

[dime10]

Regardless of what we call it, I agree it should be merged into the release.