Error when indexing into sequence dimension

[flaskeee]

Hi, I am really excited to use your tool in my project! This bug happens on an operation that is very common on language and time series models, it would be very appreciated if you can take a look at it!

Describe the bug

Error happens in the final backward pass of CROWN if a model contains indexing along the sequence dimension. i.e. new_hidden = old_hidden[:, -1, :] # (batch, seq, channel)

To Reproduce

minimal error case

import torch
from auto_LiRPA import PerturbationLpNorm, BoundedTensor, BoundedModule


class SelectLastVectorInSequence(torch.nn.Module):
    def forward(self, x):
        return x[..., -1, :]


model_width = 4
seq_len = 2
batch_size = 1
model = torch.nn.Sequential(
    torch.nn.Linear(model_width, model_width),
    torch.nn.Sigmoid(),
    torch.nn.Linear(model_width, model_width),
    SelectLastVectorInSequence(),
)
x = torch.ones((batch_size, seq_len, model_width))
ptb = PerturbationLpNorm(norm=float('inf'), x_L=-torch.ones_like(x), x_U=torch.ones_like(x))
bounded_x = BoundedTensor(x, ptb)
lirpa_model = BoundedModule(model, torch.empty_like(x))
lb, ub = lirpa_model.compute_bounds(x=(bounded_x,), method='CROWN')

Full script outputs are shown below

Traceback (most recent call last):
  File "/home/zjc/temp/bug_report/minimum_error.py", line 23, in <module>
    lb, ub = lirpa_model.compute_bounds(x=(bounded_x,), method='CROWN')
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zjc/.cache/pypoetry/virtualenvs/bug-report-jF5ToeBy-py3.11/lib/python3.11/site-packages/auto_LiRPA/bound_general.py", line 1316, in compute_bounds
    return self._compute_bounds_main(C=C,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zjc/.cache/pypoetry/virtualenvs/bug-report-jF5ToeBy-py3.11/lib/python3.11/site-packages/auto_LiRPA/bound_general.py", line 1422, in _compute_bounds_main
    ret = self.backward_general(
          ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zjc/.cache/pypoetry/virtualenvs/bug-report-jF5ToeBy-py3.11/lib/python3.11/site-packages/auto_LiRPA/backward_bound.py", line 337, in backward_general
    A, lower_b, upper_b = l.bound_backward(
                          ^^^^^^^^^^^^^^^^^
  File "/home/zjc/.cache/pypoetry/virtualenvs/bug-report-jF5ToeBy-py3.11/lib/python3.11/site-packages/auto_LiRPA/operators/add_sub.py", line 44, in bound_backward
    uA_y = _bound_oneside(last_uA, y)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zjc/.cache/pypoetry/virtualenvs/bug-report-jF5ToeBy-py3.11/lib/python3.11/site-packages/auto_LiRPA/operators/add_sub.py", line 41, in _bound_oneside
    return self.broadcast_backward(last_A, w)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zjc/.cache/pypoetry/virtualenvs/bug-report-jF5ToeBy-py3.11/lib/python3.11/site-packages/auto_LiRPA/operators/base.py", line 451, in broadcast_backward
    assert A.shape[2:] == shape[1:]  # skip the spec and batch dimension.
           ^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

Captured computation graph

Node 12 and 13 are the result of indexing op.

[BoundInput(name=/0, inputs=[], perturbed=True),
 BoundParams(name=/1, inputs=[], perturbed=False),
 BoundParams(name=/2, inputs=[], perturbed=False),
 BoundParams(name=/3, inputs=[], perturbed=False),
 BoundParams(name=/4, inputs=[], perturbed=False),
 BoundTranspose(name=/5, inputs=[/1], perturbed=False),
 BoundMatMul(name=/6, inputs=[/0, /5], perturbed=True),
 BoundAdd(name=/7, inputs=[/2, /6], perturbed=True),
 BoundSigmoid(name=/8, inputs=[/7], perturbed=True),
 BoundTranspose(name=/9, inputs=[/3], perturbed=False),
 BoundMatMul(name=/10, inputs=[/8, /9], perturbed=True),
 BoundAdd(name=/11, inputs=[/4, /10], perturbed=True),
 BoundConstant(name=/12, value=-1),
 BoundGather(name=/13, inputs=[/11, /12], perturbed=True)]

Remarks

• A.shape == (4, 1, 3, 4) and shape == (1, 2, 4)
• node under question (the target of Bound.bound_backward()) was node /11

System configuration:

• OS: Ubuntu 22.04
• Python version: 3.11
• Pytorch Version: 2.2.2
• Hardware: Error persists across GPU(RTX Titan) and CPU-only
• 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, clean Poetry environment with only auto_LiRPA and its dependencies installed, poetry config and lock file attached.

Screenshots
If applicable, add screenshots to help explain your problem.

Additional context
Add any other context about the problem here.