code: def viz_square(data, normalize=True, cmap=plt.cm.gray, padsize=1, padval=0):
      
      # normalize to 0-1 range
      if normalize:
          data -= data.min()
          data /= data.max()
      n = int(np.ceil(np.sqrt(data.shape[0]))) # force square 
      padding = ((0, n ** 2 - data.shape[0]), (0, padsize), (0, padsize)) + ((0, 0),) * (data.ndim - 3)
      data = np.pad(data, padding, mode='constant', constant_values=(padval, padval))
      # tile the filters into an image
      data = data.reshape((n, n) + data.shape[1:]).transpose((0, 2, 1, 3) + tuple(range(4, data.ndim + 1)))
      data = data.reshape((n * data.shape[1], n * data.shape[3]) + data.shape[4:])
      plt.matshow(data,cmap=cmap)
docs: takes a np.ndarray of shape (n, height, width) or (n, height, width, channels)
      visualize each (height, width) thing in a grid of size approx. sqrt(n) by sqrt(n)
      However, this only draws first input channel