code: def _batch_gather_with_broadcast(params, indices, axis):
      """Like batch_gather, but broadcasts to the left of axis."""
      # batch_gather assumes...
      #   params.shape =  [A1,...,AN, B1,...,BM]
      #   indices.shape = [A1,...,AN, C]
      # which gives output of shape
      #                   [A1,...,AN, C, B1,...,BM]
      # Here we broadcast dims of each to the left of `axis` in params, and left of
      # the rightmost dim in indices, e.g. we can
      # have
      #   params.shape =  [A1,...,AN, B1,...,BM]
      #   indices.shape = [a1,...,aN, C],
      # where ai broadcasts with Ai.
  
      # leading_bcast_shape is the broadcast of [A1,...,AN] and [a1,...,aN].
      leading_bcast_shape = ps.broadcast_shape(
          ps.shape_slice(params, np.s_[:axis]),
          ps.shape_slice(indices, np.s_[:-1]),
      )
      params = _broadcast_with(
          params,
          ps.concat(
              (leading_bcast_shape, _int32ones(ps.rank(params) - axis)),
              axis=0,
          ),
      )
      indices = _broadcast_with(
          indices,
          ps.concat((leading_bcast_shape, _int32ones(1)), axis=0),
      )
      return tf.gather(
          params,
          indices,
          batch_dims=tensorshape_util.rank(indices.shape) - 1,
      )
  
imports: import itertools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 12 for calling a global function with example usage: _batch_gather_with_broadcast(params, indices, axis) and returns: tf