alexlee-gk

import matplotlib.pyplot as plt
import numpy as np


def vis_square(data, grid_shape=None, padsize=1, padval=0, cmap=None, data_min=None, data_max=None):
  data_min = data_min if data_min is not None else data.min()
  data_max = data_max if data_max is not None else data.max()
  data = (data - data_min) / (data_max - data_min)

  lead_shape = data.shape[:-3]