Pruning of verified domains during BaB applies incompatible mask to update bounds_info.lower_bound

[ilpleutdehors]

Describe the bug
When abcrown uses branch and bound to compute bounds on subdomains, there is an issue when the program tries to apply a mask to the bounds_info.lower_bound in the method mask_batch_dim_inplace of an object of class WorkingIntermBoundsInfo. We try to apply a mask that is incompatible with the size of the bounds_info.lower_bound.

To Reproduce

1. You can find the files to reproduce the error attached to this issue. You just have to extract them in the alpha-beta-crown/complete-verifier directory.
   mask_error.zip
2. Then, you have to run them with the command python abcrown.py --config mask_error/config_mask_error.yaml. An error should appear.
3. This is the error that is raised when running the example.

Traceback (most recent call last):
  File "~/alpha-beta-CROWN/complete_verifier/abcrown.py", line 414, in <module>
    abcrown.main()
  File "~/alpha-beta-CROWN/complete_verifier/abcrown.py", line 368, in main
    verified_status = self.complete_verifier(
                      ^^^^^^^^^^^^^^^^^^^^^^^
  File "~/alpha-beta-CROWN/complete_verifier/complete_verifier_func.py", line 211, in complete_verifier
    l, num_domains_visited, ret = self.bab(
                                  ^^^^^^^^^
  File "~/alpha-beta-CROWN/complete_verifier/complete_verifier_func.py", line 142, in bab
    batch_result = general_bab(
                   ^^^^^^^^^^^^
  File "~/alpha-beta-CROWN/complete_verifier/activation_split/bab_bootstrap.py", line 297, in general_bab
    solve_packet: SolvePacket = branch_and_bound_solve(
                                ^^^^^^^^^^^^^^^^^^^^^^^
  File "~/alpha-beta-CROWN/complete_verifier/activation_split/stage_solve.py", line 79, in branch_and_bound_solve
    coreResults = update_bounds_core(
                  ^^^^^^^^^^^^^^^^^^^
  File "~/alpha-beta-CROWN/complete_verifier/activation_split/update_bounds_phases.py", line 460, in update_bounds_core
    ret = net.update_bounds_precompute_extract(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "~/alpha-beta-CROWN/complete_verifier/activation_split/decision_precompute.py", line 476, in update_bounds_precompute_extract
    WorkingIntermBoundsInfo.from_net(self, move=True).mask_batch_dim_inplace(
  File "~/alpha-beta-CROWN/complete_verifier/state/intermediate_bounds.py", line 66, in mask_batch_dim_inplace
    bounds_info.lower_bound = bounds_info.lower_bound[batch_mask]
                              ~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^
IndexError: The shape of the mask [120] at index 0 does not match the shape of the indexed tensor [73, 32] at index 0

4. The alpha-beta-CROWN repository was left unchanged.

System configuration:

• OS: Red Hat Enterprise Linux (Rocky Linux 8.10 (Green Obsidian))
• Python version: Python 3.11
• Pytorch Version: PyTorch 2.11
• Hardware: Tesla V100-PCIE-16GB and NVIDIA A100 80GB PCIe (experiments on two configs)
• 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

Additional context
The mask batch_mask should help to remove domains that were previously verified from the bounds_info.lower_bound tensor.
However, the mask is not adapted to the size of bounds_info.lower_bound, resulting in an error.
My intuition is that bounds_info.lower_bound got pruned of already verified domains beforehand. Indeed, the size of the first dimension of bounds_info.lower_bound (73) is the same as the size of batch_mask (120) minus the number of verified domains (47).
I was not able to find the source of the first pruning event.