Compilations performed to instantiate lazy constants aren't cached

[jekbradbury]

This compilation isn't cached anywhere, so code like the following:

big_lazy_x = jnp.ones((jax.device_count(), 1e6))
for _ in range(100):
  _ = pmap(lambda x: x)(big_lazy_x)

will result in one compilation per device per step of the function that instantiates big_lazy_x.
This should be fixed by the new lazy infra (#1668).

[mattjj]

By the way, that link into the code may have bitrotted, in that I don't think it's pointing to the right place anymore. (What I do to make a link is select a line and then press 'y' in GitHub's web interface to get a pointer to a specific commit. Those don't rot!)

I'm pretty sure it used to point to _instantiate_device_constant, which indeed isn't cached. And indeed #1668 fixes this! I'll add a test to that PR...