don't over-instantiate literals, fixes #1361
don't over-instantiate literals, fixes #1361

Co-authored-by: Dougal Maclaurin <0bb64e81bea11b43bfc7c06c1607c3af64f31e6d@google.com>

google/jax,google/jax,google/jax,google/jax,tensorflow/probability,tensorflow/probability