code: def _squeezed_einsum(einsum_fn, formula, *args):
      """Modified einsum that squeezes size 1 dimensions from arguments."""
  
      lhs, rhs = formula.split("->")
      lhs = lhs.split(",")
      new_args = []
      new_formula = []
  
      for indices, arg in zip(lhs, args):
          axis = []
          new_indices = ""
          for i, index in enumerate(indices):
              if arg.shape[i] == 1:
                  # This axis can be removed.
                  axis.append(i)
              else:
                  new_indices += index
          new_args.append(tf.squeeze(arg, axis=axis))
          new_formula.append(new_indices)
  
      new_formula = "{}->{}".format(",".join(new_formula), rhs)
      result = einsum_fn(new_formula, *new_args)
      return result
  
imports: 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 joint_distribution
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.experimental.marginalize.logeinsumexp import logeinsumexp
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
label_desc: Use image node_id 1 for calling a global function with example usage: _squeezed_einsum(einsum_fn, formula) and returns: result