code: def _reduce_grad(grad, shape):
      """Reduce gradient on the broadcast dimension
      If there is broadcast in forward pass, gradients need to be reduced on
      broadcast dimension. This function checks the input tensor shape and
      gradient shape and perform the reduction.
      Parameters
      ----------
      grad: Tensor
          Gradient tensor
      shape: tuple
          Shape of input tensor
      Returns
      -------
      Tensor
      """
      grad_shape = grad.shape[1:]
      in_shape = shape[1:]
      if in_shape == grad_shape:
          # no need to reduce
          return grad
      num_to_squeeze = len(grad_shape) - len(in_shape)
      # pad inshape
      in_shape = (1,) * num_to_squeeze + in_shape
      reduce_idx = np.asarray(
          np.nonzero(np.asarray(grad_shape) - np.asarray(in_shape))
      )
      reduce_idx += 1  # skip batch dim
      reduce_idx_tensor = tf.constant(
          tuple(reduce_idx.flatten().tolist()), dtype=tf.int32
      )
      grad = tf.reduce_sum(grad, axis=reduce_idx_tensor, keepdims=True)
      return tf.reshape(grad, shape)
  
imports: import numpy
  import tensorflow
  from ..._sparse_ops import _bwd_segment_cmp, _csrmask, _csrmm, _csrsum, _gsddmm, _gspmm, _scatter_add, _segment_reduce
  from ...base import ALL, is_all
  from ...heterograph_index import create_unitgraph_from_csr
  from .tensor import asnumpy, context, copy_to, tensor, zerocopy_from_numpy
label_desc: Use image node_id 3 for calling a global function with example usage: _reduce_grad(grad, shape) and returns: tf, grad