Use onp in numpy_math.astuple; previous change to avoid int(ndarray) …

…did not work.

PiperOrigin-RevId: 574584936

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

@@ -17,6 +17,7 @@
 import collections
 import functools
 import numpy as np
+import numpy as onp  # Disable JAX rewrite.  # pylint: disable=reimported
 
 from tensorflow_probability.python.internal.backend.numpy import _utils as utils
 from tensorflow_probability.python.internal.backend.numpy.numpy_array import _reverse
@@ -169,7 +170,7 @@ def _astuple(x):
   # In version 1.25 this was deprecated, causing a warning to be issued in the
   # below try/except. To avoid that, we just fall through in the case of an
   # np.ndarray.
-  if not isinstance(x, np.ndarray):
+  if not isinstance(x, onp.ndarray):
     try:
       return (int(x),)
     except TypeError: