def METHOD_NAME(a, b):
    """
    Computes the equivalent of `tf.einsum('ij,bjk->bik', a, b)`, but
    works for both dense and sparse inputs.
    :param a: Tensor or SparseTensor with rank 2.
    :param b: Tensor or SparseTensor with rank 3.
    :return: Tensor or SparseTensor with rank 3.
    """
    a_shp = tf.shape(a)
    b_shp = tf.shape(b)
    b_t = ops.transpose(b, (1, 2, 0))
    b_t = ops.reshape(b_t, tf.stack((b_shp[1], -1)))
    output = dot(a, b_t)
    output = ops.reshape(output, tf.stack((a_shp[0], b_shp[2], -1)))
    output = ops.transpose(output, (2, 0, 1))
    return output