Short-circuit _reduce_logsumexp when input is scalar, to avoid accidentally returning a [1]-shaped array.

PiperOrigin-RevId: 817734032

Author: brianwa84

tensorflow_probability/python/experimental/math/manual_special_functions_test.py

@@ -339,11 +339,11 @@ def test_patching(self, exp, log, expm1, log1p, logsumexp, softplus):
       log1p_calls += 1
       self.assertEqual(log1p_calls, log1p.call_count)
 
-      tf.math.reduce_logsumexp(0.)
+      tf.math.reduce_logsumexp([0.])
       logsumexp_calls += 1
       self.assertEqual(logsumexp_calls, logsumexp.call_count)
 
-      tf.reduce_logsumexp(0.)
+      tf.reduce_logsumexp([0.])
       logsumexp_calls += 1
       self.assertEqual(logsumexp_calls, logsumexp.call_count)
 

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

@@ -416,6 +416,8 @@ def _reduce_logsumexp(input_tensor, axis=None, keepdims=False, name=None):  # py
           or np.issubdtype(dtype, np.complexfloating)):
     # Match TF error
     raise TypeError('Input must be either real or complex')
+  if not input_tensor.shape:
+    return input_tensor
   if input_tensor.size == 0:
     # On empty arrays, mimic TF in returning `-inf` instead of failing, and
     # preserve error message if `axis` arg is incompatible with an empty array.