Shape error for 2 MaxPool layers in alpha-CROWN

[phK3]

Describe the bug
Concretization of linear bounds for alpha-CROWN fails for a model with two convolutional and two maxpool layers.

To Reproduce
Execute the following code:

import torch 
import torch.nn as nn

from auto_LiRPA import BoundedModule, BoundedTensor, PerturbationLpNorm

net = nn.Sequential(nn.Conv2d(1, 2, 2, padding=1), nn.MaxPool2d(2), nn.Conv2d(2, 2, 1, padding=1), nn.MaxPool2d(2), nn.Flatten())

x = torch.randn(1, 1, 2, 2)
model = BoundedModule(net, x) 

y = model(x)
print("concrete output: ", y)

data_min = -torch.ones(4).reshape(input_shape)
data_max = torch.ones(4).reshape(input_shape)
center = 0.5 * (data_min + data_max)
ptb = PerturbationLpNorm(x_L=data_min, x_U=data_max)
x = BoundedTensor(center, ptb)
    
print("## CROWN")
# works
lbs, ubs = model.compute_bounds(x=(x,), method='crown')
print(" CROWN finished")

print("\n## aCROWN")  
# this call fails
lbs, ubs = model.compute_bounds(x=(x,), method='alpha-crown')
print(" aCROWN finished ")

Output of the program:

concrete output:  tensor([[-0.1304,  0.6418]], grad_fn=<ViewBackward0>)
## CROWN
 CROWN finished

## aCROWN
Traceback (most recent call last):
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/maxpool_bug.py", line 58, in <module>
    lbs, ubs = model.compute_bounds(x=(x,), method='alpha-crown')
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/bound_general.py", line 1302, in compute_bounds
    ret2 = self._get_optimized_bounds(bound_side='upper', **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/optimized_bounds.py", line 475, in _get_optimized_bounds
    ret = self.compute_bounds(
          ^^^^^^^^^^^^^^^^^^^^
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/bound_general.py", line 1316, in compute_bounds
    return self._compute_bounds_main(C=C,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/bound_general.py", line 1414, in _compute_bounds_main
    self.check_prior_bounds(final)
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/bound_general.py", line 879, in check_prior_bounds
    self.check_prior_bounds(n)
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/bound_general.py", line 897, in check_prior_bounds
    self.compute_intermediate_bounds(
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/bound_general.py", line 981, in compute_intermediate_bounds
    node.lower, node.upper = self.backward_general(
                             ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/backward_bound.py", line 405, in backward_general
    lb, ub = concretize(self, batch_size, output_dim, lb, ub,
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/backward_bound.py", line 781, in concretize
    lb = lb + roots[i].perturbation.concretize(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/perturbations.py", line 253, in concretize
    return self.concretize_matrix(x, A, sign)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/philipp/VerifyNN/abCROWN_verydiff/alpha-beta-CROWN/auto_LiRPA/auto_LiRPA/perturbations.py", line 185, in concretize_matrix
    bound = A.matmul(center) + sign * A.abs().matmul(diff)
            ^^^^^^^^^^^^^^^^
RuntimeError: Expected size for first two dimensions of batch2 tensor to be: [1, 72] but got: [1, 4].

System configuration:

• OS: Ubuntu 22.04.
• Python version: 3.11.8
• Pytorch Version: 2.2.1
• Hardware: Intel Xeon E3-1200 CPU
• Have you tried to reproduce the problem in a cleanly created conda/virtualenv environment using official installation instructions and the latest code on the main branch?: Yes (bfb7997)