Don't serialize jaxpr equations into XLA metadata.

[skye]

Equations can be quite large, e.g. if it's some kind of call primitive. This can result in using many GB of host memory just to create the HLO module.

Fixes at least part of #1911.

[mattjj]

Wow, I didn't think of that, but it makes sense.

[shoyer]

Will this change effect profiling? Some of our internal tools currently show jaxprs.

[mattjj]

Yes perhaps, but the behavior here was pathological I think, potentially staging out an amount of metadata superlinear in the program size. If we need to iterate a bit to make sure the tooling still works well then we can do that, while still avoiding pathological behavior. (Also we must add tests for any behavior you want, or downstream tools need, maintained!)

[jekbradbury]

I agree (one potential middle ground would be to drop bound subjaxprs from the stored metadata); we’ll have to iterate to see what a good amount to store is (in addition to trying to make the metadata and its use in profiling tools more helpful).

Unfortunately it’s hard to test this behavior since I don’t think there’s an endpoint to access the HLO proto (or the textual HLO with metadata included) from Python right now.