docs: def get_layer_vis_square(data,
                           allow_heatmap=True,
                           normalize=True,
                           min_img_dim=100,
                           max_width=1200,
                           channel_order='RGB',
                           colormap='jet',
                           ):
      """
      Returns a vis_square for the given layer data
      Arguments:
      data -- a np.ndarray
      Keyword arguments:
      allow_heatmap -- if True, convert single channel images to heatmaps
      normalize -- whether to normalize the data when visualizing
      max_width -- maximum width for the vis_square
      """
      if channel_order not in ['RGB', 'BGR']:
          raise ValueError('Unsupported channel_order %s' % channel_order)
      if data.ndim == 1:
          # interpret as 1x1 grayscale images
          # (N, 1, 1)
          data = data[:, np.newaxis, np.newaxis]
      elif data.ndim == 2:
          # interpret as 1x1 grayscale images
          # (N, 1, 1)
          data = data.reshape((data.shape[0] * data.shape[1], 1, 1))
      elif data.ndim == 3:
          if data.shape[0] == 3:
              # interpret as a color image
              # (1, H, W,3)
              if channel_order == 'BGR':
                  data = data[[2, 1, 0], ...]  # BGR to RGB (see issue #59)
              data = data.transpose(1, 2, 0)
              data = data[np.newaxis, ...]
          else:
              # interpret as grayscale images
              # (N, H, W)
              pass
      elif data.ndim == 4:
          if data.shape[0] == 3:
              # interpret as HxW color images
              # (N, H, W, 3)
              data = data.transpose(1, 2, 3, 0)
              if channel_order == 'BGR':
                  data = data[:, :, :, [2, 1, 0]]  # BGR to RGB (see issue #59)
          elif data.shape[1] == 3:
              # interpret as HxW color images
              # (N, H, W, 3)
              data = data.transpose(0, 2, 3, 1)
              if channel_order == 'BGR':
                  data = data[:, :, :, [2, 1, 0]]  # BGR to RGB (see issue #59)
          else:
              # interpret as HxW grayscale images
              # (N, H, W)
              data = data.reshape((data.shape[0] * data.shape[1], data.shape[2], data.shape[3]))
      else:
          raise RuntimeError('unrecognized data shape: %s' % (data.shape,))
  
      return get_layer_vis_square_raw(data,
                           allow_heatmap,
                           normalize,
                           min_img_dim,
                           max_width,
                           colormap,
                           )