[jaxprs] Hoist large constants as arguments during lowering

[gnecula]

After #29882 non-scalar constants appear as core.Literal in Jaxpr.eqns.invars, and ClosedJaxpr are mostly empty of consts. While this simplifies the Jaxprs, there are several problem with the current lowering where we serialize the constants as HLO constants:

• if the constant is on a device, it gets copied to the host and is embedded in the lowered HLO. This increases memory usage, and if the constant was sharded on multiple devices it will now become replicated. Several large Google-internal tests are OOMing.
• if a constant appears in multiple places in the Jaxpr, we will make copies.
• if the constants are inlined in the HLO the XLA constant folding starts to constant-fold them, and we see warnings due to very slow constant folding.
• the constant folding changes the numerical behavior compared to current lowering strategy.

With the change in this PR, during lowering we hoist the constants out of the Jaxpr (the new function jaxpr_hoisted_consts), we pass them as arguments to the top-level function and we thread them to other functions that need them. We do yet hoist constants for AOT because we don't want to change the semantics of lower to produce a lowering and a list of constants.

I tried several places in the call stack where to hoist the constants as arguments, and update the avals, sharding, layouts for the arguments to reflect the newly added arguments. I found that it is easiest to do this fairly high up the call stack, in pxla.lower_sharding_computation because the new avals/sharding/layouts are then passed down to the lowering functions from there, but are also put in MeshComputation. I had to add several arguments to lowering functions to pass avals (they were before picked from the Jaxpr) and num_hoisted_consts. The latter makes some sense to have as an argument because it specifies the calling convention desired from the lowered functions.

This change has no effect, unless JAX_USE_SIMPLIFIED_JAXPR_CONSTANTS=True. If the flag is False (still the default) then jaxpr_hoisted_consts returns the empty list, and throughout this change num_hoisted_consts is 0. The code behavior should not change.

There are a few more tests to fix before we can flip the default for JAX_USE_SIMPLIFIED_JAXPR_CONSTANTS, to be done in several smaller follow-up PRs.