Add missing convert_to_tensor in broadcast_to for NumPy/JAX backends.

PiperOrigin-RevId: 607793016
tensorflow · Feb 16, 2024 · 6467548 · 6467548
1 parent b24cdb2
commit 6467548
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensorflow_probability/python/internal/backend/numpy/ops.py b/tensorflow_probability/python/internal/backend/numpy/ops.py
@@ -394,7 +394,8 @@ def __init__(self, *args, **kwargs):
 
 broadcast_to = utils.copy_docstring(
     'tf.broadcast_to',
-    lambda input, shape, name=None: np.broadcast_to(input, shape))
+    lambda input, shape, name=None: np.broadcast_to(
+        _convert_to_tensor(input), shape))
 
 
 def _cast(x, dtype):