docstr: Like max_pool but with input using axes ('c', 0, 1, 'b')
    (Alex Krizhevsky format)
  
  pool_shp, pool_stride and img_shp are int that represent
  the same shp in x and y.
code: def max_pool_c01b(c01b, pool_shp, pool_stride, img_shp):
      """
      Like max_pool but with input using axes ('c', 0, 1, 'b')
        (Alex Krizhevsky format)
  
      pool_shp, pool_stride and img_shp are int that represent
      the same shp in x and y.
      """
      mx = None
  
      # Compute index in pooled space of last needed pool
      # (needed = each input pixel must appear in at least one pool)
      def last_pool(im_shp, p_shp, p_strd):
          rval = int(np.ceil(float(im_shp - p_shp) / p_strd))
          assert p_strd * rval + p_shp >= im_shp
          assert p_strd * (rval - 1) + p_shp < im_shp
          return rval
  
      # Compute starting row of the last pool
      last_pool_r = (
          last_pool(img_shp, pool_shp, pool_stride) * pool_stride
      )
      # Compute number of rows needed in img for all indexes to work out
      required_r = last_pool_r + pool_shp
  
      last_pool_c = (
          last_pool(img_shp, pool_shp, pool_stride) * pool_stride
      )
      required_c = last_pool_c + pool_shp
  
      wide_infinity = pt.alloc(
          -np.inf, c01b.shape[0], required_r, required_c, c01b.shape[3]
      )
  
      c01b = set_subtensor(
          wide_infinity[:, 0:img_shp, 0:img_shp, :], c01b
      )
  
      for row_within_pool in range(pool_shp):
          row_stop = last_pool_r + row_within_pool + 1
          for col_within_pool in range(pool_shp):
              col_stop = last_pool_c + col_within_pool + 1
              cur = c01b[
                  :,
                  row_within_pool:row_stop:pool_stride,
                  col_within_pool:col_stop:pool_stride,
                  :,
              ]
              if mx is None:
                  mx = cur
              else:
                  mx = maximum(mx, cur)
      return mx
  
imports: import copy
  import numpy
  import pytest
  import pytensor.tensor
  from pytensor import shared
  from pytensor.compile.function import function
  from pytensor.compile.mode import get_default_mode, get_mode
  from pytensor.compile.ops import deep_copy_op
  from pytensor.configdefaults import config
  from pytensor.graph.basic import Apply, Variable, equal_computations
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.op import Op
  from pytensor.graph.rewriting.basic import check_stack_trace, node_rewriter, out2in
  from pytensor.graph.rewriting.utils import rewrite_graph
  from pytensor.graph.type import Type
  from pytensor.tensor.basic import alloc, as_tensor_variable
  from pytensor.tensor.elemwise import DimShuffle, Elemwise
  from pytensor.tensor.math import add, exp, maximum
  from pytensor.tensor.rewriting.basic import register_specialize
  from pytensor.tensor.rewriting.shape import ShapeFeature, local_reshape_to_dimshuffle, local_useless_reshape
  from pytensor.tensor.shape import Reshape, Shape_i, SpecifyShape, reshape, shape, specify_shape
  from pytensor.tensor.subtensor import set_subtensor
  from pytensor.tensor.type import fmatrix, iscalar, lscalar, matrix, scalar, tensor, tensor3, tensor4, vector
  from tests import unittest_tools
label_desc: Use image node_id 3 for calling the TestShapeRewriter obj's underlying member method code with example usage: obj.max_pool_c01b(c01b, pool_shp, pool_stride, img_shp) and returns: mx, rval