def bmm_maybe_select(A, B, index):
    """Slice submatrices of A by the given index and perform bmm.

    B is a 3D tensor of shape (N, D1, D2), which can be viewed as a stack of
    N matrices of shape (D1, D2). The input index is an integer vector of length M.
    A could be either:
    (1) a dense tensor of shape (M, D1),
    (2) an integer vector of length M.
    The result C is a 2D matrix of shape (M, D2)

    For case (1), C is computed by bmm:
    ::

        C[i, :] = matmul(A[i, :], B[index[i], :, :])

    For case (2), C is computed by index select:
    ::

        C[i, :] = B[index[i], A[i], :]

    Parameters
    ----------
    A : tf.Tensor
        lhs tensor
    B : tf.Tensor
        rhs tensor
    index : tf.Tensor
        index tensor

    Returns
    -------
    C : tf.Tensor
        return tensor
    """
    if A.dtype == tf.int64 and len(A.shape) == 1:
        # following is a faster version of B[index, A, :]
        B = tf.reshape(B, (-1, B.shape[2]))
        flatidx = index * B.shape[1] + A
        return tf.gather(B, flatidx)
    else:
        BB = tf.gather(B, index)
        return tf.squeeze(tf.matmul(tf.expand_dims(A, 1), BB), 1)
Slice submatrices of A by the given index and perform bmm.

    B is a 3D tensor of shape (N, D1, D2), which can be viewed as a stack of
    N matrices of shape (D1, D2). The input index is an integer vector of length M.
    A could be either:
    (1) a dense tensor of shape (M, D1),
    (2) an integer vector of length M.
    The result C is a 2D matrix of shape (M, D2)

    For case (1), C is computed by bmm:
    ::

        C[i, :] = matmul(A[i, :], B[index[i], :, :])

    For case (2), C is computed by index select:
    ::

        C[i, :] = B[index[i], A[i], :]

    Parameters
    ----------
    A : tf.Tensor
        lhs tensor
    B : tf.Tensor
        rhs tensor
    index : tf.Tensor
        index tensor

    Returns
    -------
    C : tf.Tensor
        return tensor
    
python/dgl/nn/tensorflow/utils.py