code: def test_basics_einsum(self):
      probs = np.random.rand(20) + 0.001
      probs = probs / np.sum(probs)
  
      def model():
          i = yield Root(
              categorical.Categorical(probs=probs, dtype=tf.int32)
          )
          j = yield categorical.Categorical(probs=probs, dtype=tf.int32)
          k = yield categorical.Categorical(probs=probs, dtype=tf.int32)
  
      dist = marginalize.MarginalizableJointDistributionCoroutine(model)
  
      p = tf.exp(
          dist.marginalized_log_prob(
              ["tabulate", "tabulate", "tabulate"], method="einsum"
          )
      )
      self.assertEqual(p.shape, [20, 20, 20])
      self.assertAllClose(tf.reduce_sum(p), 1.0)
  
      s = tf.exp(
          dist.marginalized_log_prob(
              ["marginalize", "marginalize", "marginalize"],
              method="einsum",
          )
      )
      self.assertAllClose(s, 1.0)
  
imports: import functools
  import itertools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import hidden_markov_model
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  import tensorflow_probability.python.experimental.marginalize
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the _MarginalizeTest obj's underlying member method code with example usage: obj.test_basics_einsum() without return types