[JAX] Fixes for tensorflow_probability for an upcoming change to jnp.array().

An upcoming change to jax.numpy.array() means that, under a transformation like jax.jit(), it will always stage its arrays into the trace. This often breaks if the array is being used for a shape calculation. Make sure we use static shapes in more places to fix test failures.

PiperOrigin-RevId: 397849299

Author: brianwa84

tensorflow_probability/python/internal/backend/meta/gen_linear_operators.py

@@ -196,6 +196,8 @@ def gen_module(module_name):
   code = code.replace('math_ops.range', 'array_ops.range')
   code = code.replace('ops.convert_to_tensor_v2_with_dispatch(',
                       'ops.convert_to_tensor(')
+  code = code.replace('ops.convert_to_tensor(dim_value)',
+                      'np.array(dim_value, np.int32)')
 
   code = code.replace('self.dtype.real_dtype', 'dtypes.real_dtype(self.dtype)')
   code = code.replace('dtype.real_dtype', 'dtypes.real_dtype(dtype)')

tensorflow_probability/python/internal/backend/numpy/gen/linear_operator.py

@@ -486,7 +486,7 @@ def _domain_dimension_tensor(self, shape=None):
     # more efficient manner, e.g. without excessive Tensor conversions.
     dim_value = tensor_shape.dimension_value(self.domain_dimension)
     if dim_value is not None:
-      return ops.convert_to_tensor(dim_value)
+      return np.array(dim_value, np.int32)
     else:
       shape = self.shape_tensor() if shape is None else shape
       return shape[-1]
@@ -530,7 +530,7 @@ def _range_dimension_tensor(self, shape=None):
     # more efficient manner, e.g. without excessive Tensor conversions.
     dim_value = tensor_shape.dimension_value(self.range_dimension)
     if dim_value is not None:
-      return ops.convert_to_tensor(dim_value)
+      return np.array(dim_value, np.int32)
     else:
       shape = self.shape_tensor() if shape is None else shape
       return shape[-2]