Array constants are generated on the host during AD

[gnecula]

I noticed that the function abstract_arrays.zeros_like_shaped_array has code like np.broadcast_to(np.array(0, aval.dtype), aval.shape) (using regular NumPy). These values will be computed on the host during tracing and then may be passed to the device. Is this the desired behavior?

For example, when we take gradients of constant functions:

x = np.ones((10, 10), dtype=np.float32)
jax.make_jaxpr(jax.grad(lambda x: 42.))(x)

we get

{ lambda a ; b.
  let 
  in (a,) }

where a 10x10 array of 0s is built during tracing and passed to the function as the array a.
Wouldn't it be better to compute the array of 0s directly on the device?

{ lambda  ; a.
   let  b = broadcast_in_dim[ broadcast_dimensions=(  )
                            shape=(10, 10) ] 0.0
  in (b,) }

Incidentally, this affects things like jax2tf where we may end up with large constants embedded in the resulting graph, when instead their computation can be embedded in the graph.

[mattjj]

This was intended at one point, pre-lazy-sublanguage #1668 and pre-omnistaging #3370, since there is special logic for handling broadcasted numpy constants efficiently in XLA lowering. (Hence why it's written with np.broadcast_to, which uses stride tricks for the broadcasted representation, unlike np.zeros!) But that special handling is only present for XLA lowering, and not jax2tf lowering.

I think we don't need this logic anymore. There might be edge cases where it's more efficient in eager mode, but I can't think of one, and they're probably not that interesting.

I agree we should replace it with jnp operations, since that will benefit jax2tf! I can take that on.

[gnecula]

Thanks, I can also do it.

[gnecula]

@mattjj I have started on this in #7102. The actual change to use lax instead of np is very small, but there are interactions with jax2tf (fixed) and with numpy 3.6.13 (not yet fixed). I can continue to work on this, just FYI, so that you don't start afresh.