def unravel_index(idx, shape):
    """Convert the flattened ind to the coordinate array

    Parameters
    ----------
    idx : int or tvm.expr.IntImm
        The 1D index

    shape : tuple of int
        Shape of the tensor

    Returns
    -------
    indices : tuple of int or tvm.expr.IntImm
        Corresponding coordinate of the 1D index
    """
    indices = []
    for i in range(len(shape) - 1, -1, -1):
        indices.append(idx % shape[i])
        idx = idx // shape[i]
    indices = indices[::-1]
    return indices