Replace einsum in JDAB test for TF 2.5 compatibility.

PiperOrigin-RevId: 374774401

Author: brianwa84

tensorflow_probability/python/distributions/BUILD

@@ -3030,7 +3030,6 @@ 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,7 +30,6 @@
 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
@@ -494,9 +493,7 @@ def test_unit_sample_shape(self):
     @tfd.JointDistributionCoroutineAutoBatched
     def dist():
       x = yield tfd.Normal(loc=tf.zeros([3]), scale=1., name='x')
-      if ps.rank(x) != 1:
-        raise ValueError('Unexpected shape.')
-      yield tfd.Bernoulli(logits=tf.reduce_sum(x), name='y')
+      yield tfd.Bernoulli(logits=tf.einsum('n->', x), name='y')
 
     for sample_shape in [(), 1, [1], [1, 1], [2]]:
       self.assertAllEqual(