[Capture] Add execution support for broadcast_in_dim_p and iota_p

[albi3ro]

Context:

While we can capture the creation of arrays with a dynamic shape like jnp.ones, jnp.zeros, and jnp.arange, jax does not support executing these equations:

def f(n):
    return jax.numpy.ones((n,))

jaxpr = jax.make_jaxpr(f)(2)
jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 4)

...
XlaRuntimeError: UNKNOWN: [/var/folders/k1/0v_kvphn55lgf_45kntf1hqm0000gq/T/ipykernel_16612/1590642475.py:2:11](http://localhost:8888/var/folders/k1/0v_kvphn55lgf_45kntf1hqm0000gq/T/ipykernel_16612/1590642475.py#line=1): error: 'mhlo.dynamic_broadcast_in_dim' op can't be translated to XLA HLO
...

I think jax.core.eval_jaxpr is trying to jit and lower the broadcast_in_dim primitive, but can't. This problem propogates to executing any jaxpr with a PlxprInterpreter

Description of the Change:

Registers special handling for broadcast_in_dim and iota.

Benefits:

We can execute the creation of dynamically shaped arrays.

Possible Drawbacks:

Sensitive to the location of jax.lax.broadcast_in_dim_p and jax.lax.iota_p. So this may break if jax moves things around.

Related GitHub Issues:

[sc-82653] [sc-82656]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Please upload report for BASE (dynamic-capture-hop-2@ab090bf). Learn more about missing BASE report.

Additional details and impacted files

@@                   Coverage Diff                    @@
##             dynamic-capture-hop-2    #6865   +/-   ##
========================================================
  Coverage                         ?   99.60%           
========================================================
  Files                            ?      477           
  Lines                            ?    45263           
  Branches                         ?        0           
========================================================
  Hits                             ?    45083           
  Misses                           ?      180           
  Partials                         ?        0           

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

[mudit2812]

Looks good. I'd suggest expanding the docstring of _fill_in_shape_with_dyn_shape, and also adding some more inline comments to the primitive registrations about the changes. I had to look at the jaxpr of a dummy function locally to convince myself why this works 😅 . Maybe it's worth including some info in the intro_to_dynamic_shapes.md file.

[mudit2812]

Looks good. I think a few more dev comments should be added, otherwise this looks approval ready to me.

[PietropaoloFrisoni]

Looks excellent, I just don't understand the examples in the docstrings. If you could clarify that it would be great 👍

[PietropaoloFrisoni]

🚀