Use jnp.zeros rather than jnp.empty in places where values are important.

Why? JAX has always lowered `jnp.empty` to `jnp.zeros`, because in early versions of XLA there was no way to create an uninitialized buffer. In a future JAX release, `jnp.empty` will lower to `jax.lax.empty`, which creates an uninitialized buffer. This is closer to the intent of `jnp.empty`, and can lead to more efficient code when the author knows that initializing the array is not required.

These problematic uses were detected by temporarily making `jnp.empty` return `NaN` (for float & complex dtypes) or the maximum value (for int and bool dtypes), and running global presubmit tests.

PiperOrigin-RevId: 900691452

Author: vanderplas

tensorflow_probability/python/internal/backend/numpy/tensor_array_ops.py

@@ -54,7 +54,7 @@ def __init__(self,
     self._dtype = utils.numpy_dtype(dtype)
     if data is None:
       if JAX_MODE and size is not None and element_shape is not None:
-        data = np.empty((size,) + tuple(element_shape), dtype=self._dtype)
+        data = np.zeros((size,) + tuple(element_shape), dtype=self._dtype)
       # Can be useful for finding failure cases in JAX TensorArray-using code.
       # elif JAX_MODE:
       #   raise ValueError(