Replace einsum in JDAB test for TF 2.5 compatibility.

PiperOrigin-RevId: 374718460

Author: davmre

tensorflow_probability/python/distributions/BUILD

@@ -3030,6 +3030,7 @@ multi_substrate_py_test(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:test_util",
         # tensorflow/compiler/jit dep,
     ],

tensorflow_probability/python/distributions/joint_distribution_auto_batched_test.py

@@ -30,6 +30,7 @@
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 
+from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import test_util
 
 tfb = tfp.bijectors
@@ -493,7 +494,9 @@ def test_unit_sample_shape(self):
     @tfd.JointDistributionCoroutineAutoBatched
     def dist():
       x = yield tfd.Normal(loc=tf.zeros([3]), scale=1., name='x')
-      yield tfd.Bernoulli(logits=tf.einsum('n->', x), name='y')
+      if ps.rank(x) != 1:
+        raise ValueError('Unexpected shape.')
+      yield tfd.Bernoulli(logits=tf.reduce_sum(x), name='y')
 
     for sample_shape in [(), 1, [1], [1, 1], [2]]:
       self.assertAllEqual(