# Repeat along any dimension
import numpy as np
def repeat(x, n, dim):
    if dim == -1:
        dim = len(x.shape) - 1
    return x.view(int(np.prod(x.shape[:dim+1])), 1, int(np.prod(x.shape[dim+1:]))).repeat(1,n,1).view(*x.shape[:dim], n * x.shape[dim], *x.shape[dim+1:])