Full Code of KaidiXu/auto_LiRPA for AI

master ca767f1d8c0a cached
249 files
1.5 MB
383.1k tokens
1697 symbols
1 requests
Download .txt
Showing preview only (1,597K chars total). Download the full file or copy to clipboard to get everything.
Repository: KaidiXu/auto_LiRPA
Branch: master
Commit: ca767f1d8c0a
Files: 249
Total size: 1.5 MB

Directory structure:
gitextract_yr8n0etx/

├── .github/
│   └── ISSUE_TEMPLATE/
│       └── bug_report.md
├── .gitignore
├── .readthedocs.yaml
├── CONTRIBUTORS
├── LICENSE
├── README.md
├── auto_LiRPA/
│   ├── __init__.py
│   ├── backward_bound.py
│   ├── beta_crown.py
│   ├── bound_general.py
│   ├── bound_multi_gpu.py
│   ├── bound_op_map.py
│   ├── bound_ops.py
│   ├── bounded_tensor.py
│   ├── concretize_bounds.py
│   ├── concretize_func.py
│   ├── cuda/
│   │   ├── cuda_kernels.cu
│   │   └── cuda_utils.cpp
│   ├── cuda_utils.py
│   ├── edit_graph.py
│   ├── eps_scheduler.py
│   ├── forward_bound.py
│   ├── interval_bound.py
│   ├── jacobian.py
│   ├── linear_bound.py
│   ├── operators/
│   │   ├── __init__.py
│   │   ├── activation_base.py
│   │   ├── activations.py
│   │   ├── add_sub.py
│   │   ├── base.py
│   │   ├── bivariate.py
│   │   ├── clampmult.py
│   │   ├── constant.py
│   │   ├── convex_concave.py
│   │   ├── convolution.py
│   │   ├── cut_ops.py
│   │   ├── dropout.py
│   │   ├── dtype.py
│   │   ├── gelu.py
│   │   ├── indexing.py
│   │   ├── jacobian.py
│   │   ├── leaf.py
│   │   ├── linear.py
│   │   ├── logical.py
│   │   ├── minmax.py
│   │   ├── normalization.py
│   │   ├── pooling.py
│   │   ├── reduce.py
│   │   ├── relu.py
│   │   ├── reshape.py
│   │   ├── resize.py
│   │   ├── rnn.py
│   │   ├── s_shaped.py
│   │   ├── shape.py
│   │   ├── slice_concat.py
│   │   ├── softmax.py
│   │   ├── solver_utils.py
│   │   ├── tile.py
│   │   └── trigonometric.py
│   ├── opt_pruner.py
│   ├── optimize_graph.py
│   ├── optimized_bounds.py
│   ├── output_constraints.py
│   ├── parse_graph.py
│   ├── patches.py
│   ├── perturbations.py
│   ├── solver_module.py
│   ├── tools.py
│   ├── utils.py
│   └── wrapper.py
├── doc/
│   ├── .gitignore
│   ├── Makefile
│   ├── README.md
│   ├── api.rst
│   ├── conf.py
│   ├── index.rst
│   └── process.py
├── examples/
│   ├── .gitignore
│   ├── __init__.py
│   ├── language/
│   │   ├── .gitignore
│   │   ├── Transformer/
│   │   │   ├── Transformer.py
│   │   │   ├── __init__.py
│   │   │   ├── modeling.py
│   │   │   └── utils.py
│   │   ├── data_utils.py
│   │   ├── language_utils.py
│   │   ├── lstm.py
│   │   ├── oracle.py
│   │   ├── preprocess/
│   │   │   ├── pre_compute_lm_scores.py
│   │   │   └── preprocess_sst.py
│   │   └── train.py
│   ├── sequence/
│   │   ├── .gitignore
│   │   ├── __init__.py
│   │   ├── data_utils.py
│   │   ├── lstm.py
│   │   └── train.py
│   ├── simple/
│   │   ├── invprop.py
│   │   ├── lp_full.py
│   │   ├── mip_lp_solver.py
│   │   ├── models/
│   │   │   └── spectral_NOR_MLP_B.pth
│   │   └── toy.py
│   └── vision/
│       ├── .gitignore
│       ├── bound_option.py
│       ├── cifar_training.py
│       ├── custom_op.py
│       ├── data/
│       │   ├── .gitignore
│       │   ├── ImageNet64/
│       │   │   └── imagenet_data_loader.py
│       │   └── tinyImageNet/
│       │       ├── .gitignore
│       │       └── tinyimagenet_download.sh
│       ├── datasets.py
│       ├── efficient_convolution.py
│       ├── imagenet_training.py
│       ├── jacobian.py
│       ├── models/
│       │   ├── __init__.py
│       │   ├── densenet.py
│       │   ├── densenet_imagenet.py
│       │   ├── densenet_no_bn.py
│       │   ├── feedforward.py
│       │   ├── mobilenet.py
│       │   ├── resnet.py
│       │   ├── resnet18.py
│       │   ├── resnext.py
│       │   ├── resnext_imagenet64.py
│       │   ├── vnncomp_resnet.py
│       │   ├── wide_resnet_cifar.py
│       │   └── wide_resnet_imagenet64.py
│       ├── pretrained/
│       │   ├── cifar_2c2f.pth
│       │   ├── kw_mnist.pth
│       │   ├── mnist_a_adv.pth
│       │   ├── mnist_cnn_small.pth
│       │   ├── mnist_fc_3layer.pth
│       │   └── test_min_max.pth
│       ├── save_intermediate_bound.py
│       ├── simple_training.py
│       ├── simple_verification.py
│       ├── tinyimagenet_training.py
│       ├── verify_two_node.py
│       └── weight_perturbation_training.py
├── setup.py
└── tests/
    ├── .gitignore
    ├── data/
    │   ├── .gitignore
    │   ├── avgpool_test_data
    │   ├── beta_crown_test_data
    │   ├── bound_ops_data
    │   ├── ckpt_lstm
    │   ├── ckpt_transformer
    │   ├── constant_test_data
    │   ├── conv1d_test_data_3-0-2
    │   ├── conv1d_test_data_3-0-3
    │   ├── conv1d_test_data_3-1-2
    │   ├── conv1d_test_data_3-1-3
    │   ├── conv1d_test_data_4-0-2
    │   ├── conv1d_test_data_4-0-3
    │   ├── conv1d_test_data_4-1-2
    │   ├── conv1d_test_data_4-1-3
    │   ├── distinct_patches_test_data
    │   ├── invprop/
    │   │   ├── ood.onnx
    │   │   ├── ood_reference
    │   │   └── simple_reference
    │   ├── jacobian_test_data
    │   ├── language_test_data
    │   ├── maxpool_test_data_3-0-3-0
    │   ├── maxpool_test_data_3-0-3-1
    │   ├── maxpool_test_data_3-1-3-0
    │   ├── maxpool_test_data_3-1-3-1
    │   ├── maxpool_test_data_4-0-4-0
    │   ├── maxpool_test_data_4-0-4-1
    │   ├── maxpool_test_data_4-1-4-0
    │   ├── maxpool_test_data_4-1-4-1
    │   ├── min_max_test_data
    │   ├── rectangle_patches_test_data
    │   ├── resnet_patches_test_data
    │   ├── s_shape_test_data
    │   ├── test_constrained_concretize
    │   ├── test_general_shape_data
    │   ├── test_perturbation_data
    │   ├── test_save_data
    │   ├── vision_clip_test_data
    │   ├── vision_test_data
    │   └── weight_perturbation_test_data
    ├── data_64/
    │   ├── avgpool_test_data
    │   ├── bound_ops_data
    │   ├── constant_test_data
    │   ├── conv1d_test_data_3-0-2
    │   ├── conv1d_test_data_3-0-3
    │   ├── conv1d_test_data_3-1-2
    │   ├── conv1d_test_data_3-1-3
    │   ├── conv1d_test_data_4-0-2
    │   ├── conv1d_test_data_4-0-3
    │   ├── conv1d_test_data_4-1-2
    │   ├── conv1d_test_data_4-1-3
    │   ├── general_shape_data
    │   ├── invprop/
    │   │   ├── ood_reference
    │   │   └── simple_reference
    │   ├── jacobian_test_data
    │   ├── maxpool_test_data_3-0-3-0
    │   ├── maxpool_test_data_3-0-3-1
    │   ├── maxpool_test_data_3-1-3-0
    │   ├── maxpool_test_data_3-1-3-1
    │   ├── maxpool_test_data_4-0-4-0
    │   ├── maxpool_test_data_4-0-4-1
    │   ├── maxpool_test_data_4-1-4-0
    │   ├── maxpool_test_data_4-1-4-1
    │   ├── min_max_test_data
    │   ├── rectangle_patches_test_data
    │   ├── resnet_patches_test_data
    │   ├── s_shape_test_data
    │   ├── test_constrained_concretize
    │   ├── test_general_shape_data
    │   ├── test_save_data
    │   ├── vision_clip_test_data
    │   ├── vision_test_data
    │   └── weight_perturbation_test_data
    ├── test_1d_activation.py
    ├── test_2d_activation.py
    ├── test_avgpool.py
    ├── test_bound_ops.py
    ├── test_branching_heuristics.py
    ├── test_clip_domains.py
    ├── test_constant.py
    ├── test_constrained_concretize.py
    ├── test_conv.py
    ├── test_conv1d.py
    ├── test_distinct_patches.py
    ├── test_examples.py
    ├── test_examples_ci.py
    ├── test_general_nonlinear.py
    ├── test_general_shape.py
    ├── test_identity.py
    ├── test_invprop.py
    ├── test_jacobian.py
    ├── test_language_models.py
    ├── test_linear_cnn_model.py
    ├── test_linear_model.py
    ├── test_maxpool.py
    ├── test_min_max.py
    ├── test_perturbation.py
    ├── test_rectangle_patches.py
    ├── test_resnet_patches.py
    ├── test_s_shaped.py
    ├── test_save_intermediate.py
    ├── test_simple_verification.py
    ├── test_state_dict_name.py
    ├── test_tensor_storage.py
    ├── test_upsample.py
    ├── test_vision_models.py
    ├── test_vision_models_hardtanh.py
    ├── test_weight_perturbation.py
    └── testcase.py

================================================
FILE CONTENTS
================================================

================================================
FILE: .github/ISSUE_TEMPLATE/bug_report.md
================================================
---
name: Bug report
about: Create a report to help us improve
title: ''
labels: ''
assignees: ''

---

**Describe the bug**
A clear and concise description of what the bug is.

**To Reproduce**
Please provide us with the following to receive timely help:
1. A minimum example to reproduce the bug. Keep your code as short as possible but still directly runnable.
2. Model files, especially when the bug is only triggered on specific models.
3. **Complete** outputs of the program when the bug is triggered. Please do **not** just include the last few lines. If it's very long, you can use [PasteBin](https://pastebin.com/) or upload to a file-sharing service.
4. Detailed instructions to reproduce the problem. If you changed part of our tool, please rebase your changes to main branch and push your changes to a fork so we can investigate easier.

Without the above information, you might not be able to receive timely help from us.


**System configuration:**
 - OS: [e.g. Ubuntu 22.04. Windows and MacOS are not supported.]
 - Python version: [e.g., Python 3.8]
 - Pytorch Version: [e.g., PyTorch 1.12]
 - Hardware: [e.g., RTX 4090]
 - 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/No]

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

**Additional context**
Add any other context about the problem here.


================================================
FILE: .gitignore
================================================
tmp
build
__pycache__
*.egg-info
dist
*.swp
*.swo
*.log
.trace_graph
Verified_ret*.npy
Verified-acc*.npy
vnn-comp_*.npz
*.tar.gz
verifier_log_*
.vscode/
*.pt
.idea
*.so
release
*.compiled
.DS_Store
*.out
*.txt
release
release_abcrown
cachier
out.csv
results.csv


================================================
FILE: .readthedocs.yaml
================================================
# Read the Docs configuration file
# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details

# Required
version: 2

# Set the version of Python and other tools you might need
build:
  os: ubuntu-20.04
  tools:
    python: "3.11"

# Build documentation in the docs/ directory with Sphinx
sphinx:
   configuration: doc/conf.py

# Optionally declare the Python requirements required to build your docs
python:
   install:
    - method: pip
      path: .
    - requirements: doc/requirements.txt

================================================
FILE: CONTRIBUTORS
================================================
Team leaders:
* Faculty: Huan Zhang (huan@huan-zhang.com), UIUC
* Student: Xiangru Zhong (xiangru4@illinois.edu), UIUC

Current developers (* indicates members of VNN-COMP 2025 team):
* \*Duo Zhou (duozhou2@illinois.edu), UIUC
* \*Keyi Shen (keyis2@illinois.edu), UIUC (graduated, now at Georgia Tech)
* \*Hesun Chen (hesunc2@illinois.edu), UIUC
* \*Haoyu Li (haoyuli5@illinois.edu), UIUC
* \*Ruize Gao (ruizeg2@illinois.edu), UIUC
* \*Hao Cheng (haoc539@illinois.edu), UIUC
* Zhouxing Shi (zhouxingshichn@gmail.com), UCLA/UC Riverside
* Lei Huang (leih5@illinois.edu), UIUC
* Taobo Liao (taobol2@illinois.edu), UIUC
* Jorge Chavez (jorgejc2@illinois.edu), UIUC

Past developers:
* Hongji Xu (hx84@duke.edu), Duke University (intern with Prof. Huan Zhang)
* Christopher Brix (brix@cs.rwth-aachen.de), RWTH Aachen University
* Hao Chen (haoc8@illinois.edu), UIUC
* Keyu Lu (keyulu2@illinois.edu), UIUC
* Kaidi Xu (kx46@drexel.edu), Drexel University
* Sanil Chawla (schawla7@illinois.edu), UIUC
* Linyi Li (linyi2@illinois.edu), UIUC
* Zhuolin Yang (zhuolin5@illinois.edu), UIUC
* Zhuowen Yuan (realzhuowen@gmail.com), UIUC
* Qirui Jin (qiruijin@umich.edu), University of Michigan
* Shiqi Wang (sw3215@columbia.edu), Columbia University
* Yihan Wang (yihanwang@ucla.edu), UCLA
* Jinqi (Kathryn) Chen (jinqic@cs.cmu.edu), CMU

auto_LiRPA is currently supported in part by the National Science Foundation (NSF; award 2331967, 2525287), the AI2050 program at Schmidt Science, the Virtual Institute for Scientific Software (VISS) at Georgia Tech, the University Research Program at Toyota Research Institute (TRI), and a Mathworks research award.

The team acknowledges the financial and advisory support from Prof. Zico Kolter (zkolter@cs.cmu.edu), Prof. Cho-Jui Hsieh (chohsieh@cs.ucla.edu), Prof. Suman Jana (suman@cs.columbia.edu), Prof. Bo Li (lbo@illinois.edu), and Prof. Xue Lin (xue.lin@northeastern.edu) during 2021 - 2023.


================================================
FILE: LICENSE
================================================
Copyright (C) 2021-2025 The α,β-CROWN Team
See CONTRIBUTORS for the list of all contributors and their affiliations.
    Team leaders: 
         Faculty: Huan Zhang <huan@huan-zhang.com> (UIUC)
         Student: Xiangru Zhong <xiangru4@illinois.edu> (UIUC)
    Current developers:
                  Duo Zhou <duozhou2@illinois.edu> (UIUC)
                  Keyi Shen <keyis2@illinois.edu> (UIUC/Georgia Tech)
                  Hesun Chen <hesunc2@illinois.edu> (UIUC)
                  Haoyu Li <haoyuli5@illinois.edu> (UIUC)
                  Ruize Gao <ruizeg2@illinois.edu> (UIUC)
                  Hao Cheng <haoc539@illinois.edu> (UIUC)
                  Zhouxing Shi <zhouxingshichn@gmail.com> (UCLA/UC Riverside)
                  Lei Huang <leih5@illinois.edu> (UIUC)
                  Taobo Liao <taobol2@illinois.edu> (UIUC)
                  Jorge Chavez <jorgejc2@illinois.edu> (UIUC)

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


================================================
FILE: README.md
================================================
# auto_LiRPA: Automatic Linear Relaxation based Perturbation Analysis for Neural Networks

[![Documentation Status](https://readthedocs.org/projects/auto-lirpa/badge/?version=latest)](https://auto-lirpa.readthedocs.io/en/latest/?badge=latest)
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](http://PaperCode.cc/AutoLiRPA-Demo)
[![Video Introduction](https://img.shields.io/badge/play-video-red.svg)](http://PaperCode.cc/AutoLiRPA-Video)
[![BSD license](https://img.shields.io/badge/License-BSD-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)

<p align="center">
<a href="http://PaperCode.cc/AutoLiRPA-Video"><img src="http://www.huan-zhang.com/images/upload/lirpa/auto_lirpa_2.png" width="45%" height="45%" float="left"></a>
<a href="http://PaperCode.cc/AutoLiRPA-Video"><img src="http://www.huan-zhang.com/images/upload/lirpa/auto_lirpa_1.png" width="45%" height="45%" float="right"></a>
</p>

## What's New?
- [α,β-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git) (using `auto_LiRPA` as its core library) is the winner of [VNN-COMP 2025](https://sites.google.com/view/vnn2025) and is **ranked top-1** in all [scored benchmarks](https://github.com/VNN-COMP/vnncomp2025_results/blob/main/SCORING-SMALL-TOL/latex/main.pdf). (08/2025)
- Bounding of computation graphs containing Jacobian operators now supports more nonlinear operators (e.g., ```tanh```, ```sigmoid```), enabling verification of [continuous-time Lyapunov stability](https://github.com/Verified-Intelligence/Two-Stage_Neural_Controller_Training). (12/2025)
- [α,β-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git) (using `auto_LiRPA` as its core library) is the winner of [VNN-COMP 2024](https://sites.google.com/view/vnn2024). Our tool is **ranked top-1** in all benchmarks (including 12 [regular track](https://github.com/ChristopherBrix/vnncomp2024_results/blob/main/SCORING/latex/results_regular_track.pdf) and 9 [extended track](https://github.com/ChristopherBrix/vnncomp2024_results/blob/main/SCORING/latex/results_extended_track.pdf) benchmarks). (08/2024)
- The [INVPROP algorithm](https://arxiv.org/pdf/2302.01404.pdf) allows to compute overapproximationsw of preimages (the set of inputs of an NN generating a given output set) and tighten bounds using output constraints. (03/2024)
- Branch-and-bound support for non-ReLU and general nonlinearities ([GenBaB](https://arxiv.org/pdf/2405.21063)) with optimizable bounds (α-CROWN) for new nonlinear functions (sin, cos, GeLU). We achieve significant improvements on verifying neural networks with non-ReLU nonlinearities such as Transformers, LSTM, and [ML4ACOPF](https://github.com/AI4OPT/ml4acopf_benchmark). (09/2023)
- [α,β-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git) ([alpha-beta-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git)) (using `auto_LiRPA` as its core library) **won** [VNN-COMP 2023](https://sites.google.com/view/vnn2023). (08/2023)
- Bound computation for higher-order computational graphs to support bounding Jacobian, Jacobian-vector products, and [local Lipschitz constants](https://arxiv.org/abs/2210.07394). (11/2022)
- Our neural network verification tool [α,β-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git) ([alpha-beta-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git)) (using `auto_LiRPA` as its core library) **won** [VNN-COMP 2022](https://sites.google.com/view/vnn2022). Our library supports the large CIFAR100, TinyImageNet and ImageNet models in VNN-COMP 2022. (09/2022)
- Implementation of **general cutting planes** ([GCP-CROWN](https://arxiv.org/pdf/2208.05740.pdf)), support of more activation functions and improved performance and scalability. (09/2022)
- Our neural network verification tool [α,β-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git) ([alpha-beta-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git)) **won** [VNN-COMP 2021](https://sites.google.com/view/vnn2021) **with the highest total score**, outperforming 11 SOTA verifiers. α,β-CROWN uses the `auto_LiRPA` library as its core bound computation library. (09/2021)
- [Optimized CROWN/LiRPA](https://arxiv.org/pdf/2011.13824.pdf) bound (α-CROWN) for ReLU, **sigmoid**, **tanh**, and **maxpool** activation functions, which can significantly outperform regular CROWN bounds. See [simple_verification.py](examples/vision/simple_verification.py#L59) for an example. (07/31/2021)
- Handle split constraints for ReLU neurons ([β-CROWN](https://arxiv.org/pdf/2103.06624.pdf)) for complete verifiers. (07/31/2021)
- A memory efficient GPU implementation of backward (CROWN) bounds for
convolutional layers. (10/31/2020)
- Certified defense models for downscaled ImageNet, TinyImageNet, CIFAR-10, LSTM/Transformer. (08/20/2020)
- Adding support to **complex vision models** including DenseNet, ResNeXt and WideResNet. (06/30/2020)
- **Loss fusion**, a technique that reduces training cost of tight LiRPA bounds
(e.g. CROWN-IBP) to the same asymptotic complexity of IBP, making LiRPA based certified
defense scalable to large datasets (e.g., TinyImageNet, downscaled ImageNet). (06/30/2020)
- **Multi-GPU** support to scale LiRPA based training to large models and datasets. (06/30/2020)
- Initial release. (02/28/2020)

## Introduction

`auto_LiRPA` is a library for automatically deriving and computing bounds with
linear relaxation based perturbation analysis (LiRPA) (e.g.
[CROWN](https://arxiv.org/pdf/1811.00866.pdf) and
[DeepPoly](https://files.sri.inf.ethz.ch/website/papers/DeepPoly.pdf)) for
neural networks, which is a useful tool for formal robustness verification. We
generalize existing LiRPA algorithms for feed-forward neural networks to a
graph algorithm on general computational graphs, defined by PyTorch.
Additionally, our implementation is also automatically **differentiable**,
allowing optimizing network parameters to shape the bounds into certain
specifications (e.g., certified defense). You can find [a video ▶️ introduction
here](http://PaperCode.cc/AutoLiRPA-Video).

Our library supports the following algorithms:

* Backward mode LiRPA bound propagation ([CROWN](https://arxiv.org/pdf/1811.00866.pdf)/[DeepPoly](https://files.sri.inf.ethz.ch/website/papers/DeepPoly.pdf))
* Backward mode LiRPA bound propagation with optimized bounds ([α-CROWN](https://arxiv.org/pdf/2011.13824.pdf))
* Backward mode LiRPA bound propagation with split constraints ([β-CROWN](https://arxiv.org/pdf/2103.06624.pdf) for ReLU, and [GenBaB](https://arxiv.org/pdf/2405.21063) for general nonlinear functions)
* Generalized backward mode LiRPA bound propagation with general cutting plane constraints ([GCP-CROWN](https://arxiv.org/pdf/2208.05740.pdf))
* Backward mode LiRPA bound propagation with bounds tightened using output constraints ([INVPROP](https://arxiv.org/pdf/2302.01404.pdf))
* Generalized backward mode LiRPA bound propagation for higher-order computational graphs  ([Shi et al., 2022](https://arxiv.org/abs/2210.07394))
* Forward mode LiRPA bound propagation ([Xu et al., 2020](https://arxiv.org/pdf/2002.12920))
* Forward mode LiRPA bound propagation with optimized bounds (similar to [α-CROWN](https://arxiv.org/pdf/2011.13824.pdf))
* Interval bound propagation ([IBP](https://arxiv.org/pdf/1810.12715.pdf))
* Hybrid approaches, e.g., Forward+Backward, IBP+Backward ([CROWN-IBP](https://arxiv.org/pdf/1906.06316.pdf)), [α,β-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git) ([alpha-beta-CROWN](https://github.com/Verified-Intelligence/alpha-beta-CROWN.git))
* MIP/LP formulation of neural networks

Our library allows automatic bound derivation and computation for general
computational graphs, in a similar manner that gradients are obtained in modern
deep learning frameworks -- users only define the computation in a forward
pass, and `auto_LiRPA` traverses through the computational graph and derives
bounds for any nodes on the graph.  With `auto_LiRPA` we free users from
deriving and implementing LiPRA for most common tasks, and they can simply
apply LiPRA as a tool for their own applications.  This is especially useful
for users who are not experts of LiRPA and cannot derive these bounds manually
(LiRPA is significantly more complicated than backpropagation).

## Technical Background in 1 Minute

Deep learning frameworks such as PyTorch represent neural networks (NN) as
a computational graph, where each mathematical operation is a node and edges
define the flow of computation:

<p align="center">
<a href="http://PaperCode.cc/AutoLiRPA-Video"><img src="http://www.huan-zhang.com/images/upload/lirpa/auto_LiRPA_background_1.png" width="80%"></a>
</p>

Normally, the inputs of a computation graph (which defines a NN) are data and
model weights, and PyTorch goes through the graph and produces model prediction
(a bunch of numbers):

<p align="center">
<a href="http://PaperCode.cc/AutoLiRPA-Video"><img src="http://www.huan-zhang.com/images/upload/lirpa/auto_LiRPA_background_2.png" width="80%"></a>
</p>

Our `auto_LiRPA` library conducts perturbation analysis on a computational
graph, where the input data and model weights are defined within some
user-defined ranges.  We get guaranteed output ranges (bounds):

<p align="center">
<a href="http://PaperCode.cc/AutoLiRPA-Video"><img src="http://www.huan-zhang.com/images/upload/lirpa/auto_LiRPA_background_3.png" width="80%"></a>
</p>

## Installation

Python 3.11+ and PyTorch 2.0+ are required.
It is highly recommended to have a pre-installed PyTorch
that matches your system and our version requirement
(see [PyTorch Get Started](https://pytorch.org/get-started)).
Then you can install `auto_LiRPA` via:

```bash
git clone https://github.com/Verified-Intelligence/auto_LiRPA
cd auto_LiRPA
pip install .
```

If you intend to modify this library, use `pip install -e .` instead.

## Quick Start

First define your computation as a `nn.Module` and wrap it using
`auto_LiRPA.BoundedModule()`. Then, you can call the `compute_bounds` function
to obtain certified lower and upper bounds under input perturbations:

```python
from auto_LiRPA import BoundedModule, BoundedTensor, PerturbationLpNorm

# Define computation as a nn.Module.
class MyModel(nn.Module):
    def forward(self, x):
        # Define your computation here.

model = MyModel()
my_input = load_a_batch_of_data()
# Wrap the model with auto_LiRPA.
model = BoundedModule(model, my_input)
# Define perturbation. Here we add Linf perturbation to input data.
ptb = PerturbationLpNorm(norm=np.inf, eps=0.1)
# Make the input a BoundedTensor with the pre-defined perturbation.
my_input = BoundedTensor(my_input, ptb)
# Regular forward propagation using BoundedTensor works as usual.
prediction = model(my_input)
# Compute LiRPA bounds using the backward mode bound propagation (CROWN).
lb, ub = model.compute_bounds(x=(my_input,), method="backward")
```

Checkout
[examples/vision/simple_verification.py](examples/vision/simple_verification.py)
for a complete but very basic example.

<a href="http://PaperCode.cc/AutoLiRPA-Demo"><img align="left" width=64 height=64 src="https://colab.research.google.com/img/colab_favicon_256px.png"></a>
We also provide a [Google Colab Demo](http://PaperCode.cc/AutoLiRPA-Demo) including an example of computing verification
bounds for a 18-layer ResNet model on CIFAR-10 dataset. Once the ResNet model
is defined as usual in Pytorch, obtaining provable output bounds is as easy as
obtaining gradients through autodiff. Bounds are efficiently computed on GPUs.

## More Working Examples

We provide [a wide range of examples](doc/src/examples.md) of using `auto_LiRPA`:

* [Basic Bound Computation on a Toy Neural Network (simplest example)](examples/simple/toy.py)
* [Basic Bound Computation with **Robustness Verification** of Neural Networks as an example](doc/src/examples.md#basic-bound-computation-and-robustness-verification-of-neural-networks)
* [MIP/LP Formulation of Neural Networks](examples/simple/mip_lp_solver.py)
* [Basic **Certified Adversarial Defense** Training](doc/src/examples.md#basic-certified-adversarial-defense-training)
* [Large-scale Certified Defense Training on **ImageNet**](doc/src/examples.md#certified-adversarial-defense-on-downscaled-imagenet-and-tinyimagenet-with-loss-fusion)
* [Certified Adversarial Defense Training on Sequence Data with **LSTM**](doc/src/examples.md#certified-adversarial-defense-training-for-lstm-on-mnist)
* [Certifiably Robust Language Classifier using **Transformers**](doc/src/examples.md#certifiably-robust-language-classifier-with-transformer-and-lstm)
* [Certified Robustness against **Model Weight Perturbations**](doc/src/examples.md#certified-robustness-against-model-weight-perturbations-and-certified-defense)
* [Bounding **Jacobian** and **local Lipschitz constants**](examples/vision/jacobian.py)
* [Compute an Overapproximate of Neural Network **Preimage**](examples/simple/invprop.py)

`auto_LiRPA` has also been used in the following works:
* [**α,β-CROWN for complete neural network verification**](https://github.com/Verified-Intelligence/alpha-beta-CROWN)
* [**Fast certified robust training**](https://github.com/shizhouxing/Fast-Certified-Robust-Training)
* [**Computing local Lipschitz constants**](https://github.com/shizhouxing/Local-Lipschitz-Constants)

## Full Documentations

For more documentations, please refer to:

* [Documentation homepage](https://auto-lirpa.readthedocs.io)
* [API documentation](https://auto-lirpa.readthedocs.io/en/latest/api.html)
* [Adding custom operators](https://auto-lirpa.readthedocs.io/en/latest/custom_op.html)
* [Guide](https://auto-lirpa.readthedocs.io/en/latest/paper.html) for reproducing [our NeurIPS 2020 paper](https://arxiv.org/abs/2002.12920)

## Publications

Please kindly cite our papers if you use the `auto_LiRPA` library. Full [BibTeX entries](doc/src/examples.md#bibtex-entries) can be found [here](doc/src/examples.md#bibtex-entries).

The general LiRPA based bound propagation algorithm was originally proposed in our paper:

* [Automatic Perturbation Analysis for Scalable Certified Robustness and Beyond](https://arxiv.org/pdf/2002.12920).
NeurIPS 2020.
Kaidi Xu\*, Zhouxing Shi\*, Huan Zhang\*, Yihan Wang, Kai-Wei Chang, Minlie Huang, Bhavya Kailkhura, Xue Lin, Cho-Jui Hsieh (\* Equal contribution)

The `auto_LiRPA` library is further extended to support:

* Optimized bounds (α-CROWN):

  [Fast and Complete: Enabling Complete Neural Network Verification with Rapid and Massively Parallel Incomplete Verifiers](https://arxiv.org/pdf/2011.13824.pdf). ICLR 2021. Kaidi Xu\*, Huan Zhang\*, Shiqi Wang, Yihan Wang, Suman Jana, Xue Lin and Cho-Jui Hsieh (\* Equal contribution).

* Split constraints (β-CROWN):

  [Beta-CROWN: Efficient Bound Propagation with Per-neuron Split Constraints for Complete and Incomplete Neural Network Verification](https://arxiv.org/pdf/2103.06624.pdf). NeurIPS 2021. Shiqi Wang\*, Huan Zhang\*, Kaidi Xu\*, Suman Jana, Xue Lin, Cho-Jui Hsieh and Zico Kolter (\* Equal contribution).

* General constraints (GCP-CROWN):

  [GCP-CROWN: General Cutting Planes for Bound-Propagation-Based Neural Network Verification](https://arxiv.org/abs/2208.05740). Huan Zhang\*, Shiqi Wang\*, Kaidi Xu\*, Linyi Li, Bo Li, Suman Jana, Cho-Jui Hsieh and Zico Kolter (\* Equal contribution).

* Higher-order computational graphs (Lipschitz constants and Jacobian):

  [Efficiently Computing Local Lipschitz Constants of Neural Networks via Bound Propagation](https://arxiv.org/abs/2210.07394). NeurIPS 2022. Zhouxing Shi, Yihan Wang, Huan Zhang, Zico Kolter, Cho-Jui Hsieh.

* Branch-and-bound for non-ReLU and general nonlinear functions (GenBaB):

  [Neural Network Verification with Branch-and-Bound for General Nonlinearities](https://arxiv.org/pdf/2405.21063). TACAS 2025. Zhouxing Shi\*, Qirui Jin\*, Zico Kolter, Suman Jana, Cho-Jui Hsieh, Huan Zhang (\* Equal contribution).

* Tightening of bounds and preimage computation using the INVPROP algorithm:

  [Provably Bounding Neural Network Preimages](https://arxiv.org/pdf/2302.01404.pdf). NeurIPS 2023. Suhas Kotha\*, Christopher Brix\*, Zico Kolter, Krishnamurthy (Dj) Dvijotham\*\*, Huan Zhang\*\* (\* Equal contribution; \*\* Equal advising).

Certified training (verification-aware training by optimizing bounds) using `auto_LiRPA` is improved with:

* Much shorter warmup schedule and faster training:

  [Fast Certified Robust Training with Short Warmup](https://arxiv.org/pdf/2103.17268.pdf). NeurIPS 2021. Zhouxing Shi\*, Yihan Wang\*, Huan Zhang, Jinfeng Yi and Cho-Jui Hsieh (\* Equal contribution).

* Training-time branch-and-bound:

  [Certified Training with Branch-and-Bound: A Case Study on Lyapunov-stable Neural Control](https://arxiv.org/abs/2411.18235). Zhouxing Shi, Cho-Jui Hsieh, and Huan Zhang.


## Developers and Copyright

Team leaders:
* Faculty: Huan Zhang (huan@huan-zhang.com), UIUC
* Student: Xiangru Zhong (xiangru4@illinois.edu), UIUC

Current developers (* indicates members of VNN-COMP 2025 team):
* \*Duo Zhou (duozhou2@illinois.edu), UIUC
* \*Keyi Shen (keyis2@illinois.edu), UIUC (graduated, now at Georgia Tech)
* \*Hesun Chen (hesunc2@illinois.edu), UIUC
* \*Haoyu Li (haoyuli5@illinois.edu), UIUC
* \*Ruize Gao (ruizeg2@illinois.edu), UIUC
* \*Hao Cheng (haoc539@illinois.edu), UIUC
* Zhouxing Shi (zhouxingshichn@gmail.com), UCLA/UC Riverside
* Lei Huang (leih5@illinois.edu), UIUC
* Taobo Liao (taobol2@illinois.edu), UIUC
* Jorge Chavez (jorgejc2@illinois.edu), UIUC

Past developers:
* Hongji Xu (hx84@duke.edu), Duke University (intern with Prof. Huan Zhang)
* Christopher Brix (brix@cs.rwth-aachen.de), RWTH Aachen University
* Hao Chen (haoc8@illinois.edu), UIUC
* Keyu Lu (keyulu2@illinois.edu), UIUC
* Kaidi Xu (kx46@drexel.edu), Drexel University
* Sanil Chawla (schawla7@illinois.edu), UIUC
* Linyi Li (linyi2@illinois.edu), UIUC
* Zhuolin Yang (zhuolin5@illinois.edu), UIUC
* Zhuowen Yuan (realzhuowen@gmail.com), UIUC
* Qirui Jin (qiruijin@umich.edu), University of Michigan
* Shiqi Wang (sw3215@columbia.edu), Columbia University
* Yihan Wang (yihanwang@ucla.edu), UCLA
* Jinqi (Kathryn) Chen (jinqic@cs.cmu.edu), CMU

`auto_LiRPA` is currently supported in part by the National Science Foundation (NSF; award 2331967, 2525287), the AI2050 program at Schmidt Science, the Virtual Institute for Scientific Software (VISS) at Georgia Tech, the University Research Program at Toyota Research Institute (TRI), and a Mathworks research award.

We thank the [commits](https://github.com/Verified-Intelligence/auto_LiRPA/commits) and [pull requests](https://github.com/Verified-Intelligence/auto_LiRPA/pulls) from community contributors.

Our library is released under the BSD 3-Clause license.


================================================
FILE: auto_LiRPA/__init__.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################
from .bound_general import BoundedModule
from .bound_multi_gpu import BoundDataParallel
from .bounded_tensor import BoundedTensor, BoundedParameter
from .perturbations import PerturbationLpNorm, PerturbationSynonym, PerturbationLinear
from .wrapper import CrossEntropyWrapper, CrossEntropyWrapperMultiInput
from .bound_op_map import register_custom_op, unregister_custom_op

__version__ = '0.7.0'


================================================
FILE: auto_LiRPA/backward_bound.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################
import os
import torch
from torch import Tensor
from collections import deque
from tqdm import tqdm
from .patches import Patches
from .utils import *
from .bound_ops import *
import warnings

from typing import TYPE_CHECKING, List
if TYPE_CHECKING:
    from .bound_general import BoundedModule


def batched_backward(self: 'BoundedModule', node, C, unstable_idx, batch_size,
                     bound_lower=True, bound_upper=True, return_A=None):
    if return_A is None: return_A = self.return_A
    output_shape = node.output_shape[1:]
    dim = int(prod(output_shape))
    if unstable_idx is None:
        unstable_idx = torch.arange(dim, device=self.device)
        dense = True
    else:
        dense = False
    unstable_size = get_unstable_size(unstable_idx)
    print(f'Batched CROWN: node {node}, unstable size {unstable_size}')
    crown_batch_size = self.bound_opts['crown_batch_size']
    auto_batch_size = AutoBatchSize(self.bound_opts['crown_batch_size'], self.device, vram_ratio=self.bound_opts['batched_crown_max_vram_ratio'])

    ret = []
    ret_A = {} # if return_A, we will store A here
    i = 0
    torch.cuda.empty_cache()
    with tqdm(total=unstable_size) as pbar:
        while i < unstable_size:
            crown_batch_size = auto_batch_size.batch_size
            if isinstance(unstable_idx, tuple):
                unstable_idx_batch = tuple(
                    u[i : i + crown_batch_size]
                    for u in unstable_idx
                )
                unstable_size_batch = len(unstable_idx_batch[0])
            else:
                unstable_idx_batch = unstable_idx[i : i + crown_batch_size]
                unstable_size_batch = len(unstable_idx_batch)
            auto_batch_size.record_actual_batch_size(unstable_size_batch)

            if node.patches_start and node.mode == "patches":
                assert C is None or C.type == 'Patches'
                C_batch = Patches(shape=[
                    unstable_size_batch, batch_size, *node.output_shape[1:-2], 1, 1],
                    identity=1, unstable_idx=unstable_idx_batch,
                    output_shape=[batch_size, *node.output_shape[1:]])
            elif C.type == 'OneHot':
                assert isinstance(node, (BoundLinear, BoundMatMul))
                C_batch = OneHotC(
                    [batch_size, unstable_size_batch, *node.output_shape[1:]],
                    self.device, unstable_idx_batch, None)
            else:
                assert C is None or C.type == 'eye'
                C_batch = torch.zeros([1, unstable_size_batch, dim], device=self.device)
                C_batch[0, torch.arange(unstable_size_batch), unstable_idx_batch] = 1.0
                C_batch = C_batch.expand(batch_size, -1, -1).view(
                    batch_size, unstable_size_batch, *output_shape)
            # overwrite return_A options to run backward general
            ori_return_A_option = self.return_A
            self.return_A = return_A

            batch_ret = self.backward_general(
                node, C_batch,
                bound_lower=bound_lower, bound_upper=bound_upper,
                average_A=False, need_A_only=False, unstable_idx=unstable_idx_batch)
            ret.append(batch_ret[:2])

            if len(batch_ret) > 2:
                # A found, we merge A
                batch_A = batch_ret[2]
                ret_A = merge_A(node, batch_A, ret_A)

            # restore return_A options
            self.return_A = ori_return_A_option

            pbar.update(unstable_size_batch)
            i += unstable_size_batch
            auto_batch_size.update()

    if bound_lower:
        lb = torch.cat([item[0].view(batch_size, -1) for item in ret], dim=1)
        if dense:
            # In this case, restore_sparse_bounds will not be called.
            # And thus we restore the shape here.
            lb = lb.reshape(batch_size, *output_shape)
    else:
        lb = None
    if bound_upper:
        ub = torch.cat([item[1].view(batch_size, -1) for item in ret], dim=1)
        if dense:
            # In this case, restore_sparse_bounds will not be called.
            # And thus we restore the shape here.
            ub = ub.reshape(batch_size, *output_shape)
    else:
        ub = None

    if return_A:
        return lb, ub, ret_A
    else:
        return lb, ub


def backward_general(
    self: 'BoundedModule',
    bound_node,
    C,
    start_backpropagation_at_node = None,
    bound_lower=True,
    bound_upper=True,
    average_A=False,
    need_A_only=False,
    unstable_idx=None,
    update_mask=None,
    apply_output_constraints_to: Optional[List[str]] = None,
    initial_As: Optional[dict] = None,
    initial_lb: Optional[torch.tensor] = None,
    initial_ub: Optional[torch.tensor] = None,
):
    use_beta_crown = self.bound_opts['optimize_bound_args']['enable_beta_crown']
    tighten_input_bounds = (
        self.bound_opts['optimize_bound_args']['tighten_input_bounds']
    )

    if self.invprop_enabled():
        self.invprop_init_infeasible_bounds(bound_node, C)
    if bound_node.are_output_constraints_activated_for_layer(apply_output_constraints_to):
        return self.backward_general_invprop(
            initial_As=initial_As, initial_lb=initial_lb, initial_ub=initial_ub,
            bound_node=bound_node, C=C,
            start_backpropagation_at_node=start_backpropagation_at_node,
            bound_lower=bound_lower, bound_upper=bound_upper,
            average_A=average_A, need_A_only=need_A_only,
            unstable_idx=unstable_idx, update_mask=update_mask
        )

    roots = self.roots()

    if start_backpropagation_at_node is None:
        # When output constraints are used, backward_general_with_output_constraint()
        # adds additional layers at the end, performs the backpropagation through these,
        # and then calls backward_general() on the output layer.
        # In this case, the layer we start from (start_backpropagation_at_node) differs
        # from the layer that should be bounded (bound_node)

        # When output constraints are not used, the bounded node is the one where
        # backpropagation starts.
        start_backpropagation_at_node = bound_node

    if self.verbose:
        logger.debug(f'Bound backward from {start_backpropagation_at_node.__class__.__name__}({start_backpropagation_at_node.name}) '
                     f'to bound {bound_node.__class__.__name__}({bound_node.name})')
        if isinstance(C, BatchedCrownC):
            logger.debug(f'  C: {C}')
        elif C is not None:
            logger.debug(f'  C: shape {C.shape}, type {type(C)}')
    _print_time = bool(os.environ.get('AUTOLIRPA_PRINT_TIME', 0))

    if isinstance(C, BatchedCrownC):
        # If C is a str, use batched CROWN. If batched CROWN is not intended to
        # be enabled, C must be a explicitly provided non-str object for this function.
        if need_A_only or average_A:
            raise ValueError(
                'Batched CROWN is not compatible with '
                f'need_A_only={need_A_only}, average_A={average_A}')
        ret = self.batched_backward(
            bound_node, C, unstable_idx,
            batch_size=roots[0].value.shape[0],
            bound_lower=bound_lower, bound_upper=bound_upper,
        )
        bound_node.lower, bound_node.upper = ret[:2]
        return ret

    for n in self.nodes():
        n.lA = n.uA = None

    degree_out = get_degrees(start_backpropagation_at_node)
    C, batch_size, output_dim, output_shape = self._preprocess_C(C, bound_node)

    if initial_As is None:
        start_backpropagation_at_node.lA = C if bound_lower else None
        start_backpropagation_at_node.uA = C if bound_upper else None
    else:
        for layer_name, (lA, uA) in initial_As.items():
            self[layer_name].lA = lA
            self[layer_name].uA = uA
        assert start_backpropagation_at_node.lA is not None or start_backpropagation_at_node.uA is not None
    if initial_lb is None:
        lb = torch.tensor(0., device=self.device)
    else:
        lb = initial_lb
    if initial_ub is None:
        ub = torch.tensor(0., device=self.device)
    else:
        ub = initial_ub

    # Save intermediate layer A matrices when required.
    A_record = {}

    queue = deque([start_backpropagation_at_node])
    while len(queue) > 0:
        l = queue.popleft()  # backward from l

        if l.name in self.root_names:
            continue

        # if all the succeeds are done, then we can turn to this node in the
        # next iteration.
        for l_pre in l.inputs:
            degree_out[l_pre.name] -= 1
            if degree_out[l_pre.name] == 0:
                queue.append(l_pre)

        # Initially, l.lA or l.uA will be set to C for this node.
        if l.lA is not None or l.uA is not None:
            if self.verbose:
                logger.debug(f'  Bound backward to {l} (out shape {l.output_shape})')
                if l.lA is not None:
                    logger.debug('    lA type %s shape %s',
                                 type(l.lA), list(l.lA.shape))
                if l.uA is not None:
                    logger.debug('    uA type %s shape %s',
                                 type(l.uA), list(l.uA.shape))

            if _print_time:
                start_time = time.time()

            self.backward_from[l.name].append(bound_node)

            if not l.perturbed:
                if not hasattr(l, 'forward_value'):
                    self.get_forward_value(l)
                lb, ub = add_constant_node(lb, ub, l)
                continue

            if l.zero_uA_mtx and l.zero_lA_mtx:
                # A matrices are all zero, no need to propagate.
                continue

            lA, uA = l.lA, l.uA
            if (l.name != start_backpropagation_at_node.name and use_beta_crown
                    and getattr(l, 'sparse_betas', None)):
                lA, uA, lbias, ubias = self.beta_crown_backward_bound(
                    l, lA, uA, start_node=start_backpropagation_at_node)
                lb = lb + lbias
                ub = ub + ubias

            if isinstance(l, BoundOptimizableActivation):
                # For other optimizable activation functions (TODO: unify with ReLU).
                if bound_node.name != self.final_node_name:
                    start_shape = bound_node.output_shape[1:]
                else:
                    start_shape = C.shape[0]
                l.preserve_mask = update_mask
            else:
                start_shape = None
            A, lower_b, upper_b = l.bound_backward(
                lA, uA, *l.inputs,
                start_node=bound_node, unstable_idx=unstable_idx,
                start_shape=start_shape)

            # After propagation through this node, we delete its lA, uA variables.
            if bound_node.name != self.final_name:
                del l.lA, l.uA
            if _print_time:
                torch.cuda.synchronize()
                time_elapsed = time.time() - start_time
                if time_elapsed > 5e-3:
                    print(l, time_elapsed)
            if lb.ndim > 0 and type(lower_b) == Tensor and self.conv_mode == 'patches':
                lb, ub, lower_b, upper_b = check_patch_biases(lb, ub, lower_b, upper_b)
            lb = lb + lower_b
            ub = ub + upper_b
            if self.return_A and self.needed_A_dict and bound_node.name in self.needed_A_dict:
                # FIXME remove [0][0] and [0][1]?
                if len(self.needed_A_dict[bound_node.name]) == 0 or l.name in self.needed_A_dict[bound_node.name]:
                    # A could be either patches (in this case we cannot transpose so directly return)
                    # or matrix (in this case we transpose)
                    A_record.update({
                        l.name: {
                            "lA": (
                                A[0][0].detach() if isinstance(A[0][0], Patches)
                                else A[0][0].transpose(0, 1).detach()
                            ) if A[0][0] is not None else None,
                            "uA": (
                                A[0][1].detach() if isinstance(A[0][1], Patches)
                                else A[0][1].transpose(0, 1).detach()
                            ) if A[0][1] is not None else None,
                            # When not used, lb or ub is tensor(0).
                            "lbias": lb.transpose(0, 1).detach() if lb.ndim > 1 else None,
                            "ubias": ub.transpose(0, 1).detach() if ub.ndim > 1 else None,
                            "unstable_idx": unstable_idx
                    }})
                # FIXME: solve conflict with the following case
                self.A_dict.update({bound_node.name: A_record})
                if need_A_only and set(self.needed_A_dict[bound_node.name]) == set(A_record.keys()):
                    # We have collected all A matrices we need. We can return now!
                    self.A_dict.update({bound_node.name: A_record})
                    # Do not concretize to save time. We just need the A matrices.
                    # return A matrix as a dict: {node_start.name: [A_lower, A_upper]}
                    return None, None, self.A_dict

            for i, l_pre in enumerate(l.inputs):
                add_bound(l, l_pre, lA=A[i][0], uA=A[i][1])

    if lb.ndim >= 2:
        lb = lb.transpose(0, 1)
    if ub.ndim >= 2:
        ub = ub.transpose(0, 1)

    # TODO merge into `concretize`
    if (self.cut_used and getattr(self, 'cut_module', None) is not None
            and self.cut_module.x_coeffs is not None):
        # propagate input neuron in cut constraints
        roots[0].lA, roots[0].uA = self.cut_module.input_cut(
            bound_node, roots[0].lA, roots[0].uA, roots[0].lower.size()[1:], unstable_idx,
            batch_mask=update_mask)

    lb, ub = self.concretize_bounds(
        node=bound_node,
        lower=lb,
        upper=ub,
        concretize_mode='backward',
        batch_size=batch_size,
        output_dim=output_dim,
        average_A=average_A,
        clip_neuron_selection_value=self.clip_neuron_selection_value,
        clip_neuron_selection_type=self.clip_neuron_selection_type
    )

    if self.return_A and self.needed_A_dict and bound_node.name in self.needed_A_dict:
        save_root_A(
            bound_node, A_record, self.A_dict, roots,
            self.needed_A_dict[bound_node.name],
            lb=lb, ub=ub, unstable_idx=unstable_idx)
    for root in self.roots():
        # These are saved for `save_root_A`. We do not need them afterwards.
        root.lb = root.ub = None

    if tighten_input_bounds and isinstance(bound_node, BoundInput):
        shape = bound_node.perturbation.x_L.shape
        lb_reshaped = lb.reshape(shape)
        bound_node.perturbation.x_L = lb_reshaped - lb_reshaped.detach() + torch.max(bound_node.perturbation.x_L.detach(), lb_reshaped.detach())
        ub_reshaped = ub.reshape(shape)
        bound_node.perturbation.x_U = ub_reshaped - ub_reshaped.detach() + torch.min(bound_node.perturbation.x_U.detach(), ub_reshaped.detach())

    lb = lb.view(batch_size, *output_shape) if bound_lower else None
    ub = ub.view(batch_size, *output_shape) if bound_upper else None

    # TODO merge into `concretize`
    if (self.cut_used and getattr(self, "cut_module", None) is not None
            and self.cut_module.cut_bias is not None):
        # propagate cut bias in cut constraints
        lb, ub = self.cut_module.bias_cut(bound_node, lb, ub, unstable_idx, batch_mask=update_mask)
        if lb is not None and ub is not None and ((lb-ub)>0).sum().item() > 0:
            # make sure there is no bug for cut constraints propagation
            print(f"Warning: lb is larger than ub with diff: {(lb-ub)[(lb-ub)>0].max().item()}")

    if self.verbose:
        logger.debug('')

    if self.invprop_enabled():
        lb, ub = self.invprop_check_infeasible_bounds(lb, ub)

    if self.return_A:
        if self.bound_opts['clip_in_alpha_crown'] and self.final_name in self.A_dict.keys():
            for v in self.A_dict[self.final_name].values():
                if v["lA"] is not None:
                    self.constraints_optimized = (v["lA"], v["lbias"])
        return lb, ub, self.A_dict
    else:
        return lb, ub


def get_unstable_size(unstable_idx):
    if isinstance(unstable_idx, tuple):
        return unstable_idx[0].numel()
    else:
        return unstable_idx.numel()


def check_optimized_variable_sparsity(self: 'BoundedModule', node):
    alpha_sparsity = None  # unknown, optimizable variables are not created for this node.
    for relu in self.relus:
        # FIXME: this hardcoded for ReLUs. Need to support other optimized nonlinear functions.
        # alpha_lookup_idx is only created for sparse-spec alphas.
        if relu.alpha_lookup_idx is not None and node.name in relu.alpha_lookup_idx:
            if relu.alpha_lookup_idx[node.name] is not None:
                # This node was created with sparse alpha
                alpha_sparsity = True
            elif self.bound_opts['optimize_bound_args']['use_shared_alpha']:
                # Shared alpha, the spec dimension is 1, and sparsity can be supported.
                alpha_sparsity = True
            else:
                alpha_sparsity = False
            break
    return alpha_sparsity


def get_sparse_C(self: 'BoundedModule', node, ref_intermediate):
    (sparse_intermediate_bounds,
     ref_intermediate_lb, ref_intermediate_ub) = ref_intermediate
    sparse_conv_intermediate_bounds = self.bound_opts.get('sparse_conv_intermediate_bounds', False)
    minimum_sparsity = self.bound_opts.get('minimum_sparsity', 0.9)
    crown_batch_size = self.bound_opts.get('crown_batch_size', 1e9)
    dim = int(prod(node.output_shape[1:]))
    batch_size = self.batch_size

    reduced_dim = False  # Only partial neurons (unstable neurons) are bounded.
    unstable_idx = None
    unstable_size = np.inf
    newC = None

    alpha_is_sparse = self.check_optimized_variable_sparsity(node)

    # NOTE: batched CROWN is so far only supported for some of the cases below

    # FIXME: C matrix shape incorrect for BoundParams.
    if (isinstance(node, BoundLinear) or isinstance(node, BoundMatMul)) and int(
            os.environ.get('AUTOLIRPA_USE_FULL_C', 0)) == 0:
        if sparse_intermediate_bounds:
            # If we are doing bound refinement and reference bounds are given,
            # we only refine unstable neurons.
            # Also, if we are checking against LP solver we will refine all
            # neurons and do not use this optimization.
            # For each batch element, we find the unstable neurons.
            unstable_idx, unstable_size = self.get_unstable_locations(
                ref_intermediate_lb, ref_intermediate_ub)
            if unstable_size == 0:
                # Do nothing, no bounds will be computed.
                reduced_dim = True
                unstable_idx = []
            elif unstable_size > crown_batch_size:
                # Create C in batched CROWN
                newC = BatchedCrownC('OneHot')
                reduced_dim = True
            elif (((0 < unstable_size <= minimum_sparsity * dim and alpha_is_sparse is None) or
                   alpha_is_sparse) and
                   len(node.output_shape) <= 2):
                # When we already have sparse alpha for this layer, we always
                # use sparse C. Otherwise we determine it by sparsity.
                # Create an abstract C matrix, the unstable_idx are the non-zero
                # elements in specifications for all batches.
                # Shouldn't use OneHotC if the output is not a 1-d tensor.
                newC = OneHotC(
                    [batch_size, unstable_size, *node.output_shape[1:]],
                    self.device, unstable_idx, None)
                reduced_dim = True
            else:
                unstable_idx = None
                del ref_intermediate_lb, ref_intermediate_ub
        if not reduced_dim:
            if dim > crown_batch_size:
                newC = BatchedCrownC('eye')
            else:
                newC = eyeC([batch_size, dim, *node.output_shape[1:]], self.device)
    elif node.patches_start and node.mode == "patches":
        if sparse_intermediate_bounds:
            unstable_idx, unstable_size = self.get_unstable_locations(
                ref_intermediate_lb, ref_intermediate_ub, conv=True)
            if unstable_size == 0:
                # Do nothing, no bounds will be computed.
                reduced_dim = True
                unstable_idx = []
            elif unstable_size > crown_batch_size:
                # Create C in batched CROWN
                newC = BatchedCrownC('Patches')
                reduced_dim = True
            # We sum over the channel direction, so need to multiply that.
            elif (sparse_conv_intermediate_bounds
                  and unstable_size <= minimum_sparsity * dim
                  and alpha_is_sparse is None) or alpha_is_sparse:
                # When we already have sparse alpha for this layer, we always
                # use sparse C. Otherwise we determine it by sparsity.
                # Create an abstract C matrix, the unstable_idx are the non-zero
                # elements in specifications for all batches.
                # The shape of patches is [unstable_size, batch, C, H, W].
                newC = Patches(
                    shape=[unstable_size, batch_size, *node.output_shape[1:-2],
                           1, 1],
                    identity=1, unstable_idx=unstable_idx,
                    output_shape=[batch_size, *node.output_shape[1:]])
                reduced_dim = True
            else:
                unstable_idx = None
                del ref_intermediate_lb, ref_intermediate_ub
        # Here we create an Identity Patches object
        if not reduced_dim:
            newC = Patches(
                None, 1, 0, [node.output_shape[1], batch_size, *node.output_shape[2:],
                *node.output_shape[1:-2], 1, 1], 1,
                output_shape=[batch_size, *node.output_shape[1:]])
    elif (isinstance(node, (BoundAdd, BoundSub)) and node.mode == "patches"
        and len(node.output_shape) >= 4):
        # FIXME: BoundAdd does not always have patches. Need to use a better way
        # to determine patches mode.
        # FIXME: We should not hardcode BoundAdd here!
        if sparse_intermediate_bounds:
            if crown_batch_size < 1e9:
                warnings.warn('Batched CROWN is not supported in this case')
            unstable_idx, unstable_size = self.get_unstable_locations(
                ref_intermediate_lb, ref_intermediate_ub, conv=True)
            if unstable_size == 0:
                # Do nothing, no bounds will be computed.
                reduced_dim = True
                unstable_idx = []
            elif (sparse_conv_intermediate_bounds
                  and unstable_size <= minimum_sparsity * dim
                  and alpha_is_sparse is None) or alpha_is_sparse:
                # When we already have sparse alpha for this layer, we always
                # use sparse C. Otherwise we determine it by sparsity.
                num_channel = node.output_shape[-3]
                # Identity patch size: (ouc_c, 1, 1, 1, out_c, 1, 1).
                patches = (
                    torch.eye(num_channel, device=self.device,
                    dtype=list(self.parameters())[0].dtype)).view(
                        num_channel, 1, 1, 1, num_channel, 1, 1)
                # Expand to (out_c, 1, unstable_size, out_c, 1, 1).
                patches = patches.expand(-1, 1, node.output_shape[-2],
                                         node.output_shape[-1], -1, 1, 1)
                patches = patches[unstable_idx[0], :,
                                  unstable_idx[1], unstable_idx[2]]
                # Expand with the batch dimension. Final shape
                # (unstable_size, batch_size, out_c, 1, 1).
                patches = patches.expand(-1, batch_size, -1, -1, -1)
                newC = Patches(
                    patches, 1, 0, patches.shape, unstable_idx=unstable_idx,
                    output_shape=[batch_size, *node.output_shape[1:]])
                reduced_dim = True
            else:
                unstable_idx = None
                del ref_intermediate_lb, ref_intermediate_ub
        if not reduced_dim:
            num_channel = node.output_shape[-3]
            # Identity patch size: (ouc_c, 1, 1, 1, out_c, 1, 1).
            patches = (
                torch.eye(num_channel, device=self.device,
                dtype=list(self.parameters())[0].dtype)).view(
                    num_channel, 1, 1, 1, num_channel, 1, 1)
            # Expand to (out_c, batch, out_h, out_w, out_c, 1, 1).
            patches = patches.expand(-1, batch_size, node.output_shape[-2],
                                     node.output_shape[-1], -1, 1, 1)
            newC = Patches(patches, 1, 0, patches.shape, output_shape=[
                batch_size, *node.output_shape[1:]])
    else:
        if sparse_intermediate_bounds:
            unstable_idx, unstable_size = self.get_unstable_locations(
                ref_intermediate_lb, ref_intermediate_ub)
            if unstable_size == 0:
                # Do nothing, no bounds will be computed.
                reduced_dim = True
                unstable_idx = []
            elif unstable_size > crown_batch_size:
                # Create in C in batched CROWN
                newC = BatchedCrownC('eye')
                reduced_dim = True
            elif (unstable_size <= minimum_sparsity * dim
                  and alpha_is_sparse is None) or alpha_is_sparse:
                newC = torch.zeros([1, unstable_size, dim], device=self.device)
                # Fill the corresponding elements to 1.0
                newC[0, torch.arange(unstable_size), unstable_idx] = 1.0
                newC = newC.expand(batch_size, -1, -1).view(
                    batch_size, unstable_size, *node.output_shape[1:])
                reduced_dim = True
            else:
                unstable_idx = None
                del ref_intermediate_lb, ref_intermediate_ub
        if not reduced_dim:
            if dim > 1000:
                warnings.warn(
                    f"Creating an identity matrix with size {dim}x{dim} for node {node}. "
                    "This may indicate poor performance for bound computation. "
                    "If you see this message on a small network please submit "
                    "a bug report.", stacklevel=2)
            if dim > crown_batch_size:
                newC = BatchedCrownC('eye')
            else:
                newC = torch.eye(dim, device=self.device).unsqueeze(0).expand(
                    batch_size, -1, -1
                ).view(batch_size, dim, *node.output_shape[1:])

    return newC, reduced_dim, unstable_idx, unstable_size


def restore_sparse_bounds(self: 'BoundedModule', node, unstable_idx,
                          unstable_size, ref_intermediate,
                          new_lower=None, new_upper=None):
    ref_intermediate_lb, ref_intermediate_ub = ref_intermediate[1:]
    batch_size = self.batch_size
    if unstable_size == 0:
        # No unstable neurons. Skip the update.
        node.lower = ref_intermediate_lb.detach().clone()
        node.upper = ref_intermediate_ub.detach().clone()
    else:
        if new_lower is None:
            new_lower = node.lower
        if new_upper is None:
            new_upper = node.upper
        # If we only calculated unstable neurons, we need to scatter the results back based on reference bounds.
        if isinstance(unstable_idx, tuple):
            lower = ref_intermediate_lb.detach().clone()
            upper = ref_intermediate_ub.detach().clone()
            # Conv layer with patches, the unstable_idx is a 3-element tuple for 3 indices (C, H,W) of unstable neurons.
            if len(unstable_idx) == 3:
                lower[:, unstable_idx[0], unstable_idx[1], unstable_idx[2]] = new_lower
                upper[:, unstable_idx[0], unstable_idx[1], unstable_idx[2]] = new_upper
            elif len(unstable_idx) == 4:
                lower[:, unstable_idx[0], unstable_idx[1], unstable_idx[2], unstable_idx[3]] = new_lower
                upper[:, unstable_idx[0], unstable_idx[1], unstable_idx[2], unstable_idx[3]] = new_upper
        else:
            # Other layers.
            lower = ref_intermediate_lb.detach().clone().reshape(batch_size, -1)
            upper = ref_intermediate_ub.detach().clone().reshape(batch_size, -1)
            lower[:, unstable_idx] = new_lower.view(batch_size, -1)
            upper[:, unstable_idx] = new_upper.view(batch_size, -1)
        node.lower = lower.view(batch_size, *node.output_shape[1:])
        node.upper = upper.view(batch_size, *node.output_shape[1:])


def get_degrees(node_start):
    if not isinstance(node_start, list):
        node_start = [node_start]
    degrees = {}
    added = {}
    queue = deque()
    for node in node_start:
        queue.append(node)
        added[node.name] = True
    while len(queue) > 0:
        l = queue.popleft()
        for l_pre in l.inputs:
            degrees[l_pre.name] = degrees.get(l_pre.name, 0) + 1
            if not added.get(l_pre.name, False):
                queue.append(l_pre)
                added[l_pre.name] = True
    return degrees


def _preprocess_C(self: 'BoundedModule', C, node):
    if isinstance(C, Patches):
        if C.unstable_idx is None:
            # Patches have size (out_c, batch, out_h, out_w, c, h, w).
            if len(C.shape) == 7:
                out_c, batch_size, out_h, out_w = C.shape[:4]
                output_dim = out_c * out_h * out_w
            else:
                out_dim, batch_size, out_c, out_h, out_w = C.shape[:5]
                output_dim = out_dim * out_c * out_h * out_w
        else:
            # Patches have size (unstable_size, batch, c, h, w).
            output_dim, batch_size = C.shape[:2]
    else:
        batch_size, output_dim = C.shape[:2]

    # The C matrix specified by the user has shape (batch, spec)
    # but internally we have (spec, batch) format.
    if not isinstance(C, (eyeC, Patches, OneHotC)):
        C = C.transpose(0, 1).reshape(
            output_dim, batch_size, *node.output_shape[1:])
    elif isinstance(C, eyeC):
        C = C._replace(shape=(C.shape[1], C.shape[0], *C.shape[2:]))
    elif isinstance(C, OneHotC):
        C = C._replace(
            shape=(C.shape[1], C.shape[0], *C.shape[2:]),
            index=C.index.transpose(0,-1),
            coeffs=None if C.coeffs is None else C.coeffs.transpose(0,-1))

    if isinstance(C, Patches) and C.unstable_idx is not None:
        # Sparse patches; the output shape is (unstable_size, ).
        output_shape = [C.shape[0]]
    elif prod(node.output_shape[1:]) != output_dim and not isinstance(C, Patches):
        # For the output node, the shape of the bound follows C
        # instead of the original output shape
        #
        # TODO Maybe don't set node.lower and node.upper in this case?
        # Currently some codes still depend on node.lower and node.upper
        output_shape = [-1]
    else:
        # Generally, the shape of the bounds match the output shape of the node
        output_shape = node.output_shape[1:]

    return C, batch_size, output_dim, output_shape


def addA(A1, A2):
    """ Add two A (each of them is either Tensor or Patches) """
    if type(A1) == type(A2):
        return A1 + A2
    elif type(A1) == Patches:
        return A1 + A2
    elif type(A2) == Patches:
        return A2 + A1
    else:
        raise NotImplementedError(f'Unsupported types for A1 ({type(A1)}) and A2 ({type(A2)}')


def add_bound(node, node_pre, lA=None, uA=None):
    """
    Propagate lA and uA to a preceding node.
    @param node:        The current bounded node
    @param node_pre:    An input of the current bounded node that needs lA, lbias ,etc. back propagated to it
    @param lA:          lA matrix associated with the current bounded node
    @param uA:          uA matrix associated with the current bounded node
    @return:
    """

    if lA is not None:
        if node_pre.lA is None:
            # First A added to this node.
            node_pre.zero_lA_mtx = node.zero_backward_coeffs_l
            node_pre.lA = lA
        else:
            node_pre.zero_lA_mtx = node_pre.zero_lA_mtx and node.zero_backward_coeffs_l
            new_node_lA = addA(node_pre.lA, lA)
            node_pre.lA = new_node_lA
    if uA is not None:
        if node_pre.uA is None:
            # First A added to this node.
            node_pre.zero_uA_mtx = node_pre.zero_backward_coeffs_u
            node_pre.uA = uA
        else:
            node_pre.zero_uA_mtx = node_pre.zero_uA_mtx and node.zero_backward_coeffs_u
            node_pre.uA = addA(node_pre.uA, uA)


def add_constant_node(lb, ub, node):
    new_lb = node.get_bias(node.lA, node.forward_value)
    new_ub = node.get_bias(node.uA, node.forward_value)
    if isinstance(lb, Tensor) and isinstance(new_lb, Tensor) and lb.ndim > 0 and lb.ndim != new_lb.ndim:
        new_lb = new_lb.reshape(lb.shape)
    if isinstance(ub, Tensor) and isinstance(new_ub, Tensor) and ub.ndim > 0 and ub.ndim != new_ub.ndim:
        new_ub = new_ub.reshape(ub.shape)
    lb = lb + new_lb # FIXME (09/16): shape for the bias of BoundConstant.
    ub = ub + new_ub
    return lb, ub


def save_root_A(node, A_record, A_dict, roots, needed_A_dict, lb, ub,
                unstable_idx):
    root_A_record = {}
    for i in range(len(roots)):
        if roots[i].lA is None and roots[i].uA is None:
            continue
        if roots[i].name in needed_A_dict:
            if roots[i].lA is not None:
                if isinstance(roots[i].lA, Patches):
                    _lA = roots[i].lA.detach()
                else:
                    _lA = roots[i].lA.transpose(0, 1).detach()
            else:
                _lA = None
            if roots[i].uA is not None:
                if isinstance(roots[i].uA, Patches):
                    _uA = roots[i].uA.detach()
                else:
                    _uA = roots[i].uA.transpose(0, 1).detach()
            else:
                _uA = None

            # Include all the bias terms except the one concretized from the
            # current root node.
            lb_ = lb - roots[i].lb if (roots[i].lb is not None) else lb
            ub_ = ub - roots[i].ub if (roots[i].ub is not None) else ub

            root_A_record.update({roots[i].name: {
                "lA": _lA,
                "uA": _uA,
                # When not used, lb or ub is tensor(0). They have been transposed above.
                "lbias": lb_.detach() if lb_.ndim > 1 else None,
                "ubias": ub_.detach() if ub_.ndim > 1 else None,
                "unstable_idx": unstable_idx
            }})

    root_A_record.update(A_record)  # merge to existing A_record
    A_dict.update({node.name: root_A_record})


def select_unstable_idx(ref_intermediate_lb, ref_intermediate_ub, unstable_locs, max_crown_size):
    """When there are too many unstable neurons, only bound those
    with the loosest reference bounds."""
    gap = (
        ref_intermediate_ub[:, unstable_locs]
        - ref_intermediate_lb[:, unstable_locs]).sum(dim=0)
    indices = torch.argsort(gap, descending=True)
    indices_selected = indices[:max_crown_size]
    indices_selected, _ = torch.sort(indices_selected)
    print(f'{len(indices_selected)}/{len(indices)} unstable neurons selected for CROWN')
    return indices_selected


def get_unstable_locations(self: 'BoundedModule', ref_intermediate_lb,
                           ref_intermediate_ub, conv=False, channel_only=False):
    # FIXME (2023): This function should be a member class of the Bound object, since the
    # definition of unstable neurons depends on the activation function.
    max_crown_size = self.bound_opts.get('max_crown_size', int(1e9))
    # For conv layer we only check the case where all neurons are active/inactive.
    unstable_masks = torch.logical_and(ref_intermediate_lb < 0, ref_intermediate_ub > 0)
    # For simplicity, merge unstable locations for all elements in this batch. TODO: use individual unstable mask.
    # It has shape (H, W) indicating if a neuron is unstable/stable.
    # TODO: so far we merge over the batch dimension to allow easier implementation.
    if channel_only:
        # Only keep channels with unstable neurons. Used for initializing alpha.
        unstable_locs = unstable_masks.sum(dim=(0,2,3)).bool()
        # Shape is consistent with linear layers: a list of unstable neuron channels (no batch dim).
        unstable_idx = unstable_locs.nonzero().squeeze(1)
    else:
        if not conv and unstable_masks.ndim > 2:
            # Flatten the conv layer shape.
            unstable_masks = unstable_masks.reshape(unstable_masks.size(0), -1)
            ref_intermediate_lb = ref_intermediate_lb.reshape(ref_intermediate_lb.size(0), -1)
            ref_intermediate_ub = ref_intermediate_ub.reshape(ref_intermediate_ub.size(0), -1)
        unstable_locs = unstable_masks.sum(dim=0).bool()
        if conv:
            # Now converting it to indices for these unstable nuerons.
            # These are locations (i,j) of unstable neurons.
            unstable_idx = unstable_locs.nonzero(as_tuple=True)
        else:
            unstable_idx = unstable_locs.nonzero().squeeze(1)

    unstable_size = get_unstable_size(unstable_idx)
    if unstable_size > max_crown_size:
        indices_seleted = select_unstable_idx(
            ref_intermediate_lb, ref_intermediate_ub, unstable_locs, max_crown_size)
        if isinstance(unstable_idx, tuple):
            unstable_idx = tuple(u[indices_seleted] for u in unstable_idx)
        else:
            unstable_idx = unstable_idx[indices_seleted]
    unstable_size = get_unstable_size(unstable_idx)

    return unstable_idx, unstable_size


def get_alpha_crown_start_nodes(
        self: 'BoundedModule',
        node,
        c=None,
        share_alphas=False,
        final_node_name=None,
    ):
    """
    Given a layer "node", return a list of following nodes after this node whose bounds
    will propagate through this node. Each element in the list is a tuple with 3 elements:
    (following_node_name, following_node_shape, unstable_idx)
    """
    # When use_full_conv_alpha is True, conv layers do not share alpha.
    sparse_intermediate_bounds = self.bound_opts.get('sparse_intermediate_bounds', False)
    use_full_conv_alpha_thresh = self.bound_opts.get('use_full_conv_alpha_thresh', 512)

    start_nodes = []

    for nj in self.backward_from[node.name]:  # Pre-activation layers.
        unstable_idx = None
        use_sparse_conv = None  # Whether a sparse-spec alpha is used for a conv output node. None for non-conv output node.
        use_full_conv_alpha = self.bound_opts.get('use_full_conv_alpha', False)

        # Find the indices of unstable neuron, used for create sparse-feature alpha.
        if (sparse_intermediate_bounds
                and isinstance(node, BoundOptimizableActivation)
                and nj.name != final_node_name and not share_alphas):
            # Create sparse optimization variables for intermediate neurons.
            # These are called "sparse-spec" alpha because we only create alpha only for
            # the intermediate of final output nodes whose bounds are needed.
            # "sparse-spec" alpha makes sense only for piece-wise linear functions.
            # For other intermediate nodes, there is no "unstable" or "stable" neuron.
            # FIXME: whether an layer has unstable/stable neurons should be in Bound obj.
            # FIXME: get_unstable_locations should be a member class of ReLU.
            if len(nj.output_name) == 1 and isinstance(self[nj.output_name[0]], (BoundRelu, BoundSignMerge, BoundMaxPool)):
                if ((isinstance(nj, (BoundLinear, BoundMatMul)))
                        and int(os.environ.get('AUTOLIRPA_USE_FULL_C', 0)) == 0):
                    # unstable_idx has shape [neuron_size_of_nj]. Batch dimension is reduced.
                    unstable_idx, _ = self.get_unstable_locations(nj.lower, nj.upper)
                elif isinstance(nj, (BoundConv, BoundAdd, BoundSub, BoundBatchNormalization)) and nj.mode == 'patches':
                    if nj.name in node.patch_size:
                        # unstable_idx has shape [channel_size_of_nj]. Batch and spatial dimensions are reduced.
                        unstable_idx, _ = self.get_unstable_locations(
                            nj.lower, nj.upper, channel_only=not use_full_conv_alpha, conv=True)
                        use_sparse_conv = False  # alpha is shared among channels. Sparse-spec alpha in hw dimension not used.
                        if use_full_conv_alpha and unstable_idx[0].size(0) > use_full_conv_alpha_thresh:
                            # Too many unstable neurons. Using shared alpha per channel.
                            unstable_idx, _ = self.get_unstable_locations(
                                nj.lower, nj.upper, channel_only=True, conv=True)
                            use_full_conv_alpha = False
                    else:
                        # Matrix mode for conv layers. Although the bound propagation started with patches mode,
                        # when A matrix is propagated to this layer, it might become a dense matrix since patches
                        # can be come very large after many layers. In this case,
                        # unstable_idx has shape [c_out * h_out * w_out]. Batch dimension is reduced.
                        unstable_idx, _ = self.get_unstable_locations(nj.lower, nj.upper)
                        use_sparse_conv = True  # alpha is not shared among channels, and is sparse in spec dimension.
            else:
                # FIXME: we should not check for fixed names here. Need to enable patches mode more generally.
                if isinstance(nj, (BoundConv, BoundAdd, BoundSub, BoundBatchNormalization)) and nj.mode == 'patches':
                    use_sparse_conv = False  # Sparse-spec alpha can never be used, because it is not a ReLU activation.

        if nj.name == final_node_name:
            # Final layer, always the number of specs as the shape.
            size_final = self[final_node_name].output_shape[1:] if c is None else c.size(1)
            # The 4-th element indicates that this start node is the final node,
            # which may be utilized by operators that do not know the name of
            # the final node.
            start_nodes.append((final_node_name, size_final, None, True))
            continue

        if share_alphas:
            # all intermediate neurons from the same layer share the same set of alphas.
            output_shape = 1
        elif isinstance(node, BoundOptimizableActivation) and node.patch_size and nj.name in node.patch_size:
            # Patches mode. Use output channel size as the spec size. This still shares some alpha, but better than no sharing.
            if use_full_conv_alpha:
                # alphas not shared among channels, so the spec dim shape is c,h,w
                # The patch size is [out_ch, batch, out_h, out_w, in_ch, H, W]. We use out_ch as the output shape.
                output_shape = node.patch_size[nj.name][0], node.patch_size[nj.name][2], node.patch_size[nj.name][3]
            else:
                # The spec dim is c only, and is shared among h, w.
                output_shape = node.patch_size[nj.name][0]
            assert not sparse_intermediate_bounds or use_sparse_conv is False  # Double check our assumption holds. If this fails, then we created wrong shapes for alpha.
        else:
            # Output is linear layer (use_sparse_conv = None), or patch converted to matrix (use_sparse_conv = True).
            assert not sparse_intermediate_bounds or use_sparse_conv is not False  # Double check our assumption holds. If this fails, then we created wrong shapes for alpha.
            output_shape = nj.lower.shape[1:]  # FIXME: for non-relu activations it's still expecting a prod.
        start_nodes.append((nj.name, output_shape, unstable_idx, False))

    return start_nodes


def merge_A(node, batch_A, ret_A):
    for key0 in batch_A:
        if key0 not in ret_A: ret_A[key0] = {}
        for key1 in batch_A[key0]:
            value = batch_A[key0][key1]
            if key1 not in ret_A[key0]:
                # create:
                ret_A[key0].update({
                    key1: {
                        "lA": value["lA"],
                        "uA": value["uA"],
                        "lbias": value["lbias"],
                        "ubias": value["ubias"],
                        "unstable_idx": value["unstable_idx"]
                    }
                })
            elif key0 == node.name:
                # merge:
                # the batch splitting only happens for current node, i.e.,
                # for other nodes the returned lA should be the same across different batches
                # so no need to repeatly merge them
                exist = ret_A[key0][key1]

                if exist["unstable_idx"] is not None:
                    if isinstance(exist["unstable_idx"], torch.Tensor):
                        merged_unstable = torch.cat([
                            exist["unstable_idx"],
                            value['unstable_idx']], dim=0)
                    elif isinstance(exist["unstable_idx"], tuple):
                        if exist["unstable_idx"]:
                            merged_unstable = tuple([
                                torch.cat([exist["unstable_idx"][idx],
                                           value['unstable_idx'][idx]], dim=0)
                                for idx in range(len(exist['unstable_idx']))]
                            )
                        else:
                            merged_unstable = None
                    else:
                        raise NotImplementedError(
                            f'Unsupported type {type(exist["unstable_idx"])}')
                else:
                    merged_unstable = None
                merge_dict = {"unstable_idx": merged_unstable}
                for name in ["lA", "uA"]:
                    if exist[name] is not None:
                        if isinstance(exist[name], torch.Tensor):
                            # for matrix the spec dim is 1
                            merge_dict[name] = torch.cat([exist[name], value[name]], dim=1)
                        else:
                            assert isinstance(exist[name], Patches)
                            # for patches the spec dim`is 0
                            merge_dict[name] = exist[name].create_similar(
                                torch.cat([exist[name].patches, value[name].patches], dim=0),
                                unstable_idx=merged_unstable
                            )
                    else:
                        merge_dict[name] = None
                for name in ["lbias", "ubias"]:
                    if exist[name] is not None:
                        # for bias the spec dim in 1
                        merge_dict[name] = torch.cat([exist[name], value[name]], dim=1)
                    else:
                        merge_dict[name] = None
                ret_A[key0][key1] = merge_dict
    return ret_A


================================================
FILE: auto_LiRPA/beta_crown.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################
from collections import OrderedDict
import numpy as np
import torch
from torch import Tensor
from .patches import Patches, inplace_unfold

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from .bound_general import BoundedModule


class SparseBeta:
    def __init__(self, shape, bias=False, betas=None, device='cpu'):
        self.device = device
        self.val = torch.zeros(shape)
        self.loc = torch.zeros(shape, dtype=torch.long, device=device)
        self.sign = torch.zeros(shape, device=device)
        self.bias = torch.zeros(shape, device=device) if bias else None
        if betas:
            for bi in range(len(betas)):
                if betas[bi] is not None:
                    self.val[bi, :len(betas[bi])] = betas[bi]
        self.val = self.val.detach().to(
            device, non_blocking=True).requires_grad_()

    def apply_splits(self, history, key):
        loc_numpy = np.zeros(self.loc.shape, dtype=np.int32)
        sign_numpy = np.zeros(self.sign.shape)
        if self.bias is not None:
            bias_numpy = np.zeros(self.bias.shape)
        for bi in range(len(history)):
            # Add history splits. (layer, neuron) is the current decision.
            split_locs, split_coeffs = history[bi][key][:2]
            split_len = len(split_locs)
            if split_len > 0:
                sign_numpy[bi, :split_len] = split_coeffs
                loc_numpy[bi, :split_len] = split_locs
                if self.bias is not None:
                    split_bias = history[bi][key][2]
                    bias_numpy[bi, :split_len] = split_bias
        self.loc.copy_(torch.from_numpy(loc_numpy), non_blocking=True)
        self.sign.copy_(torch.from_numpy(sign_numpy), non_blocking=True)
        if self.bias is not None:
            self.bias.copy_(torch.from_numpy(bias_numpy), non_blocking=True)

def get_split_nodes(self: 'BoundedModule'):
    self.split_nodes = []
    self.split_activations = {}
    splittable_activations = self.get_splittable_activations()
    self._set_used_nodes(self[self.final_name])
    for layer in self.layers_requiring_bounds:
        split_activations_ = []
        for activation_name in layer.output_name:
            activation = self[activation_name]
            if activation in splittable_activations:
                split_activations_.append(
                    (activation, activation.inputs.index(layer)))
        if split_activations_:
            if layer.lower is None and layer.upper is None:
                continue
            self.split_nodes.append(layer)
            self.split_activations[layer.name] = split_activations_
    return self.split_nodes, self.split_activations


def set_beta(self: 'BoundedModule', enable_opt_interm_bounds, parameters,
             lr_beta, lr_cut_beta, cutter, dense_coeffs_mask):
    """
    Set betas, best_betas, coeffs, dense_coeffs_mask, best_coeffs, biases
    and best_biases.
    """
    coeffs = None
    betas = []
    best_betas = OrderedDict()

    # TODO compute only once
    self.nodes_with_beta = []
    for node in self.split_nodes:
        if not hasattr(node, 'sparse_betas'):
            continue
        self.nodes_with_beta.append(node)
        if enable_opt_interm_bounds:
            for sparse_beta in node.sparse_betas.values():
                if sparse_beta is not None:
                    betas.append(sparse_beta.val)
            best_betas[node.name] = {
                beta_m: sparse_beta.val.detach().clone()
                for beta_m, sparse_beta in node.sparse_betas.items()
            }
        else:
            betas.append(node.sparse_betas[0].val)
            best_betas[node.name] = node.sparse_betas[0].val.detach().clone()

    # Beta has shape (batch, max_splits_per_layer)
    parameters.append({
        'params': [item for item in betas if item.numel() > 0],
        'lr': lr_beta, 'batch_dim': 0})

    if self.cut_used:
        self.set_beta_cuts(parameters, lr_cut_beta, betas, best_betas, cutter)

    return betas, best_betas, coeffs, dense_coeffs_mask


def set_beta_cuts(self: 'BoundedModule', parameters, lr_cut_beta, betas,
                  best_betas, cutter):
    # also need to optimize cut betas
    parameters.append({'params': self.cut_beta_params,
                        'lr': lr_cut_beta, 'batch_dim': 0})
    betas += self.cut_beta_params
    best_betas['cut'] = [beta.detach().clone() for beta in self.cut_beta_params]
    if getattr(cutter, 'opt', False):
        parameters.append(cutter.get_parameters())


def reset_beta(self: 'BoundedModule', node, shape, betas, bias=False,
               start_nodes=None):
    # Create only the non-zero beta. For each layer, it is padded to maximal length.
    # We create tensors on CPU first, and they will be transferred to GPU after initialized.
    if self.bound_opts.get('enable_opt_interm_bounds', False):
        node.sparse_betas = {
            key: SparseBeta(
                shape,
                betas=[(betas[j][i] if betas[j] is not None else None)
                        for j in range(len(betas))],
                device=self.device, bias=bias,
            ) for i, key in enumerate(start_nodes)
        }
    else:
        node.sparse_betas = [SparseBeta(
            shape, betas=betas, device=self.device, bias=bias)]


def beta_crown_backward_bound(self: 'BoundedModule', node, lA, uA, start_node=None):
    """Update A and bias with Beta-CROWN.

    Must be explicitly called at the end of "bound_backward".
    """
    # Regular Beta CROWN with single neuron split
    # Each split constraint only has single neuron (e.g., second ReLU neuron > 0).
    A = lA if lA is not None else uA
    lbias = ubias = 0

    def _bias_unsupported():
        raise NotImplementedError('Bias for beta not supported in this case.')

    if type(A) is Patches:
        if not self.bound_opts.get('enable_opt_interm_bounds', False):
            raise NotImplementedError('Sparse beta not supported in the patches mode')
        if node.sparse_betas[start_node.name].bias is not None:
            _bias_unsupported()
        # expand sparse_beta to full beta
        beta_values = (node.sparse_betas[start_node.name].val
                       * node.sparse_betas[start_node.name].sign)
        beta_indices = node.sparse_betas[start_node.name].loc
        node.masked_beta = torch.zeros(2, *node.shape).reshape(2, -1).to(A.patches.dtype)
        node.non_deter_scatter_add(
            node.masked_beta, dim=1, index=beta_indices,
            src=beta_values.to(node.masked_beta.dtype))
        node.masked_beta = node.masked_beta.reshape(2, *node.shape)
        # unfold the beta as patches, size (batch, out_h, out_w, in_c, H, W)
        A_patches = A.patches
        masked_beta_unfolded = inplace_unfold(
            node.masked_beta, kernel_size=A_patches.shape[-2:],
            padding=A.padding, stride=A.stride,
            inserted_zeros=A.inserted_zeros, output_padding=A.output_padding)
        if A.unstable_idx is not None:
            masked_beta_unfolded = masked_beta_unfolded.permute(1, 2, 0, 3, 4, 5)
            # After selection, the shape is (unstable_size, batch, in_c, H, W).
            masked_beta_unfolded = masked_beta_unfolded[A.unstable_idx[1], A.unstable_idx[2]]
        else:
            # Add the spec (out_c) dimension.
            masked_beta_unfolded = masked_beta_unfolded.unsqueeze(0)
        if node.alpha_beta_update_mask is not None:
            masked_beta_unfolded = masked_beta_unfolded[node.alpha_beta_update_mask]
        if uA is not None:
            uA = uA.create_similar(uA.patches + masked_beta_unfolded)
        if lA is not None:
            lA = lA.create_similar(lA.patches - masked_beta_unfolded)
    elif type(A) is Tensor:
        if self.bound_opts.get('enable_opt_interm_bounds', False):
            if node.sparse_betas[start_node.name].bias is not None:
                _bias_unsupported()
            # For matrix mode, beta is sparse.
            beta_values = (
                node.sparse_betas[start_node.name].val
                * node.sparse_betas[start_node.name].sign
            ).expand(A.size(0), -1, -1)
            # node.single_beta_loc has shape [batch, max_single_split].
            # Need to expand at the specs dimension.
            beta_indices = (node.sparse_betas[start_node.name].loc
                            .unsqueeze(0).expand(A.size(0), -1, -1))
            beta_bias = node.sparse_betas[start_node.name].bias
        else:
            # For matrix mode, beta is sparse.
            beta_values = (
                node.sparse_betas[0].val * node.sparse_betas[0].sign
            ).expand(A.size(0), -1, -1)
            # self.single_beta_loc has shape [batch, max_single_split].
            # Need to expand at the specs dimension.
            beta_indices = node.sparse_betas[0].loc.unsqueeze(0).expand(A.size(0), -1, -1)
            beta_bias = node.sparse_betas[0].bias
        # For conv layer, the last dimension is flattened in indices.
        beta_values = beta_values.to(A.dtype)
        if beta_bias is not None:
            beta_bias = beta_bias.expand(A.size(0), -1, -1)
        if node.alpha_beta_update_mask is not None:
            beta_indices = beta_indices[:, node.alpha_beta_update_mask]
            beta_values = beta_values[:, node.alpha_beta_update_mask]
            if beta_bias is not None:
                beta_bias = beta_bias[:, node.alpha_beta_update_mask]
        if uA is not None:
            uA = node.non_deter_scatter_add(
                uA.reshape(uA.size(0), uA.size(1), -1), dim=2,
                index=beta_indices, src=beta_values).view(uA.size())
        if lA is not None:
            lA = node.non_deter_scatter_add(
                lA.reshape(lA.size(0), lA.size(1), -1), dim=2,
                index=beta_indices, src=beta_values.neg()).view(lA.size())
        if beta_bias is not None:
            bias = (beta_values * beta_bias).sum(dim=-1)
            lbias = bias
            ubias = -bias
    else:
        raise RuntimeError(f"Unknown type {type(A)} for A")

    return lA, uA, lbias, ubias


def print_optimized_beta(acts):
    masked_betas = []
    for model in acts:
        masked_betas.append(model.masked_beta)
        if model.history_beta_used:
            print(f'{model.name} history beta', model.new_history_beta.squeeze())
        if model.split_beta_used:
            print(f'{model.name} split beta:', model.split_beta.view(-1))
            print(f'{model.name} bias:', model.split_bias)


================================================
FILE: auto_LiRPA/bound_general.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################

import copy
from typing import List
import numpy as np
import warnings
from collections import OrderedDict, deque

import torch
from torch.nn import Parameter

from .bound_op_map import bound_op_map
from .bound_ops import *
from .bounded_tensor import BoundedTensor, BoundedParameter
from .parse_graph import parse_module
from .perturbations import *
from .utils import *
from .patches import Patches
from .optimized_bounds import default_optimize_bound_args


warnings.simplefilter('once')


class BoundedModule(nn.Module):
    """Bounded module with support for automatically computing bounds.

    Args:
        model (nn.Module): The original model to be wrapped by BoundedModule.

        global_input (tuple): A dummy input to the original model. The shape of
        the dummy input should be consistent with the actual input to the model
        except for the batch dimension.

        bound_opts (dict): Options for bounds. See
        `Bound Options <bound_opts.html>`_.

        device (str or torch.device): Device of the bounded module.
        If 'auto', the device will be automatically inferred from the device of
        parameters in the original model or the dummy input.

        custom_ops (dict): A dictionary of custom operators.
        The dictionary maps operator names to their corresponding bound classes
        (subclasses of `Bound`).

    """
    def __init__(self, model, global_input, bound_opts=None,
                device='auto', verbose=False, custom_ops=None):
        super().__init__()
        if isinstance(model, BoundedModule):
            for key in model.__dict__.keys():
                setattr(self, key, getattr(model, key))
            return

        self.ori_training = model.training

        if bound_opts is None:
            bound_opts = {}
        # Default options.
        default_bound_opts = {
            'conv_mode': 'patches',
            'sparse_intermediate_bounds': True,
            'sparse_conv_intermediate_bounds': True,
            'sparse_intermediate_bounds_with_ibp': True,
            'sparse_features_alpha': True,
            'sparse_spec_alpha': True,
            'minimum_sparsity': 0.9,
            'enable_opt_interm_bounds': False,
            'crown_batch_size': np.inf,
            'forward_refinement': False,
            'forward_max_dim': int(1e9),
            # Do not share alpha for conv layers.
            'use_full_conv_alpha': True,
            'disabled_optimization': [],
            # Threshold for number of unstable neurons for each layer to disable
            #  use_full_conv_alpha.
            'use_full_conv_alpha_thresh': 512,
            'verbosity': 1 if verbose else 0,
            'optimize_graph': {'optimizer': None},
            'compare_crown_with_ibp': False,
            # Whether run an additional forward pass before computing bounds.
            'forward_before_compute_bounds': False,
            'clip_in_alpha_crown': False,
            # Whether to compute bounds for every node in the graph.
            # (rather than only the nodes whose intermediate bounds are needed.)
            'bound_every_node': False,
        }
        default_bound_opts.update(bound_opts)
        self.bound_opts = default_bound_opts
        optimize_bound_args = copy.deepcopy(default_optimize_bound_args)
        optimize_bound_args.update(
            self.bound_opts.get('optimize_bound_args', {}))
        self.bound_opts.update({'optimize_bound_args': optimize_bound_args})

        self.verbose = verbose
        self.custom_ops = custom_ops if custom_ops is not None else {}
        if device == 'auto':
            try:
                self.device = next(model.parameters()).device
            except StopIteration:
                # Model has no parameters. We use the device of input tensor.
                if isinstance(global_input, torch.Tensor):
                    self.device = global_input.device
                elif isinstance(global_input, tuple):
                    self.device = global_input[0].device
                else:
                    raise NotImplementedError( # pylint: disable=raise-missing-from
                        'Unable to decide the device. Consider providing a '
                        '`device` argument to `BoundedModule` explicitly.')
        else:
            self.device = device

        self.global_input = tuple(unpack_inputs(global_input, device=self.device))
        self.check_incompatible_nodes(model)

        self.conv_mode = self.bound_opts.get('conv_mode', 'patches')
        # Cached IBP results which may be reused
        self.ibp_lower, self.ibp_upper = None, None

        self.optimizable_activations = []
        self.relus = []  # save relu layers for convenience
        self.layers_with_constraint = []

        state_dict_copy = copy.deepcopy(model.state_dict())
        object.__setattr__(self, 'ori_state_dict', state_dict_copy)
        model.to(self.device)
        output = model(*self.global_input)
        if not isinstance(output, torch.Tensor):
            raise TypeError(
                'Output of the model is expected to be a single torch.Tensor. '
                f'Actual type: {type(output)}')
        self.final_shape = output.shape
        self.bound_opts.update({'final_shape': self.final_shape})
        self._convert(model, self.global_input)
        self._optimize_graph()
        # Compute forward_value and mark perturbed nodes
        self.forward(*self.global_input)
        self._expand_jacobian()
        self._check_patches_mode()

        self.next_split_hint = []  # Split hints, used in beta optimization.
        # Beta values for all intermediate bounds.
        # Set to None (not used) by default.
        self.best_intermediate_betas = None
        # Initialization value for intermediate betas.
        self.init_intermediate_betas = None
        # whether using cut
        self.cut_used = False
        # a placeholder for cut timestamp, which would be a non-positive int
        self.cut_timestamp = -1
        # a placeholder to save the latest samplewise mask for
        # pruning-in-iteration optimization
        self.last_update_preserve_mask = None
        # If output constraints are used, it is possible that none of the possible
        # inputs satisfy them. In this case, the lower bounds will be set to +inf,
        # and the upper bounds to -inf.
        self.infeasible_bounds = None
        self.solver_model = None
        # Needed for output constraints - the output layer should not use them
        self.final_node().is_final_node = True
        self.dynamic = False
        # This is the topk ratio for half-naive, half-constrained concretization.
        # Please check for concretize_bounds.py for more details.
        self.clip_neuron_selection_type = 'ratio'
        self.clip_neuron_selection_value = -1.0
        # A boolean tensor with shape (batchsize, ). It indicates if a batch is
        # infeasible when concretizing with constraints.
        # Always call `init_infeasible_bounds_constraints` function to initialize it.
        self.infeasible_bounds_constraints = None

        # This is designed for clipping during alpha-CROWN.
        # For each alpha-CROWN optimization iteration, the lA and lbias of the final layer
        #   will be set as `constraints_optimized` for the next iteration.
        # Please check backward_bound.py and optimized_bound for more info.
        self.constraints_optimized = None

    def nodes(self) -> List[Bound]:
        return self._modules.values()

    def get_enabled_opt_act(self):
        # Optimizable activations that are actually used and perturbed
        return [
            n for n in self.optimizable_activations
            if n.used and n.perturbed and not getattr(n, 'is_linear_op', False)
        ]

    def get_optimizable_activations(self):
        for node in self.nodes():
            if (isinstance(node, BoundOptimizableActivation)
                    and node.optimizable
                    and len(getattr(node, 'requires_input_bounds', [])) > 0
                    and node not in self.optimizable_activations):
                disabled = False
                for item in self.bound_opts.get('disable_optimization', []):
                    if item.lower() in str(type(node)).lower():
                        disabled = True
                if disabled:
                    logging.debug('Disabled optimization for %s', node)
                    continue
                if node not in self.optimizable_activations:
                    self.optimizable_activations.append(node)

    def get_perturbed_optimizable_activations(self):
        return [n for n in self.optimizable_activations if n.perturbed]

    def get_splittable_activations(self):
        """Activation functions that can be split during branch and bound."""
        return [n for n in self.nodes() if n.perturbed and n.splittable and n.used]

    def get_layers_requiring_bounds(self):
        """Layer names whose intermediate layer bounds are required."""
        intermediate_layers = []
        tighten_input_bounds = (
            self.bound_opts['optimize_bound_args']['tighten_input_bounds']
        )
        directly_optimize_layer_names = (
            self.bound_opts['optimize_bound_args']['directly_optimize']
        )
        for node in self.nodes():
            if node.name in directly_optimize_layer_names:
                intermediate_layers.append(node)
            if not node.used or not node.perturbed:
                continue
            for i in getattr(node, 'requires_input_bounds', []):
                input_node = node.inputs[i]
                if (input_node not in intermediate_layers
                        and input_node.perturbed):
                    # If not perturbed, it may not have the batch dimension.
                    # So we do not include it, and it is unnecessary.
                    intermediate_layers.append(input_node)
            if (
                node.name in self.layers_with_constraint
                or (isinstance(node, BoundInput) and tighten_input_bounds)
            ):
                if node not in intermediate_layers:
                    intermediate_layers.append(node)
        return intermediate_layers

    def check_incompatible_nodes(self, model):
        """Check whether the model has incompatible nodes that the conversion
        may be inaccurate"""
        node_types = [type(m) for m in list(model.modules())]

        if (torch.nn.Dropout in node_types
                and torch.nn.BatchNorm1d in node_types
                and any(input.shape[0] == 1 for input in self.global_input)):
            # In fact, we just need the input that is involved in the
            # dropout layer to have batch size larger than 1, but we don't know
            # which of them is, so we just check all of them.
            print('We cannot support torch.nn.Dropout and torch.nn.BatchNorm1d '
                  'at the same time!')
            print('Suggest to use another dummy input which has batch size '
                  'larger than 1 and set model to train() mode.')
            return

        if not self.ori_training and torch.nn.Dropout in node_types:
            print('Dropout operation CANNOT be parsed during conversion when '
                  'the model is in eval() mode!')
            print('Set model to train() mode!')
            self.ori_training = True

        if self.ori_training and torch.nn.BatchNorm1d in node_types:
            print('BatchNorm1d may raise error during conversion when the model'
                  ' is in train() mode!')
            print('Set model to eval() mode!')
            self.ori_training = False

    def non_deter_wrapper(self, op, *args, **kwargs):
        """Some operations are non-deterministic and deterministic mode will
        fail. So we temporary disable it."""
        if self.bound_opts.get('deterministic', False):
            torch.use_deterministic_algorithms(False)
        ret = op(*args, **kwargs)
        if self.bound_opts.get('deterministic', False):
            torch.use_deterministic_algorithms(True)
        return ret

    def non_deter_scatter_add(self, *args, **kwargs):
        return self.non_deter_wrapper(torch.scatter_add, *args, **kwargs)

    def non_deter_index_select(self, *args, **kwargs):
        return self.non_deter_wrapper(torch.index_select, *args, **kwargs)

    def set_bound_opts(self, new_opts):
        for k, v in new_opts.items():
            # assert v is not dict, 'only support change optimize_bound_args'
            if type(v) == dict:
                self.bound_opts[k].update(v)
            else:
                self.bound_opts[k] = v

    def set_gcp_relu_indicators(self, relu_layer_name, relu_indicators):
        """
        Sets the GCP (Generalized Cutting Plane) relu indicators for
        the specified ReLU layer by name.
        Args:
            relu_layer_name (str):
                The name of the ReLU layer to update.
            relu_indicators (torch.Tensor):
                A tensor containing unstable relu indices or masks.
        """
        # Search for the layer by name
        for m in self.relus:
            if m.name == relu_layer_name:
                # Set the indicators for the found ReLU layer
                m.gcp_unstable_relu_indicators = relu_indicators
                return
        # If not found, raise an error
        raise ValueError(f'No ReLU layer found with name {relu_layer_name}')

    @staticmethod
    def _get_A_norm(A):
        if not isinstance(A, (list, tuple)):
            A = (A, )
        norms = []
        for aa in A:
            if aa is not None:
                if isinstance(aa, Patches):
                    aa = aa.patches
                norms.append(aa.abs().sum().item())
            else:
                norms.append(None)
        return norms

    def __call__(self, *input, **kwargs):
        if 'method_opt' in kwargs:
            opt = kwargs['method_opt']
            kwargs.pop('method_opt')
        else:
            opt = 'forward'
        for kwarg in [
            'disable_multi_gpu', 'no_replicas', 'get_property',
            'node_class', 'att_name']:
            if kwarg in kwargs:
                kwargs.pop(kwarg)
        if opt == 'compute_bounds':
            return self.compute_bounds(**kwargs)
        else:
            return self.forward(*input, **kwargs)

    def register_parameter(self, name, param):
        r"""Adds a parameter to the module.

        The parameter can be accessed as an attribute using given name.

        Args:
            name (string): name of the parameter. The parameter can be accessed
                from this module using the given name
            param (Parameter): parameter to be added to the module.
        """
        if '_parameters' not in self.__dict__:
            raise AttributeError(
                'cannot assign parameter before Module.__init__() call')
        elif not isinstance(name, str):
            raise TypeError('parameter name should be a string. '
                            f'Got {torch.typename(name)}')
        elif name == '':
            raise KeyError('parameter name can\'t be empty string')
        elif hasattr(self, name) and name not in self._parameters:
            raise KeyError(f'attribute "{name}" already exists')

        if param is None:
            self._parameters[name] = None
        elif not isinstance(param, Parameter):
            raise TypeError(
                f'cannot assign "{torch.typename(param)}" object to '
                f'parameter "{name}" '
                '(torch.nn.Parameter or None required)')
        elif param.grad_fn:
            raise ValueError(
                f'Cannot assign non-leaf Tensor to parameter "{name}". Model '
                'parameters must be created explicitly. To express "{name}" '
                'as a function of another Tensor, compute the value in '
                'the forward() method.')
        else:
            self._parameters[name] = param

    def _named_members(self,
                       get_members_fn,
                       prefix='',
                       recurse=True,
                       remove_duplicate: bool = True,
                       **kwargs):  # pylint: disable=unused-argument
        r"""Helper method for yielding various names + members of modules."""
        memo = set()
        modules = self.named_modules(prefix=prefix) if recurse else [
                                     (prefix, self)]
        for module_prefix, module in modules:
            members = get_members_fn(module)
            for k, v in members:
                if v is None or v in memo:
                    continue
                if remove_duplicate:
                    memo.add(v)
                name = module_prefix + ('.' if module_prefix else '') + k
                # translate name to ori_name
                if name in self.node_name_map:
                    name = self.node_name_map[name]
                yield name, v

    def train(self, mode=True):
        super().train(mode)
        for node in self.nodes():
            node.train(mode=mode)

    def eval(self):
        super().eval()
        for node in self.nodes():
            node.eval()

    def to(self, *args, **kwargs):
        # Moves and/or casts some attributes except pytorch will do by default.
        for node in self.nodes():
            for attr in ['lower', 'upper', 'forward_value', 'd', 'lA',]:
                if hasattr(node, attr):
                    this_attr = getattr(node, attr)
                    if isinstance(this_attr, torch.Tensor):
                        this_attr = this_attr.to(*args, **kwargs)
                        setattr(node, attr, this_attr)

            if hasattr(node, 'interval'):
                # construct new interval
                this_attr = getattr(node, 'interval')
                setattr(node, 'interval', (this_attr[0].to(
                    *args, **kwargs), this_attr[1].to(*args, **kwargs)))

        return super().to(*args, **kwargs)

    def __getitem__(self, name):
        module = self._modules[name]
        # We never create modules that are None, the assert fixes type hints
        assert module is not None
        return module

    def roots(self):
        return [self[name] for name in self.root_names]

    def final_node(self):
        return self[self.final_name]

    def get_forward_value(self, node):
        """ Recursively get `forward_value` for `node` and its parent nodes"""
        if getattr(node, 'forward_value', None) is not None:
            return node.forward_value
        inputs = [self.get_forward_value(inp) for inp in node.inputs]
        for inp in node.inputs:
            node.from_input = node.from_input or inp.from_input
        node.input_shape = inputs[0].shape if len(inputs) > 0 else None
        fv = node.forward(*inputs)
        if isinstance(fv, (torch.Size, tuple)):
            fv = torch.tensor(fv, device=self.device)
        node.forward_value = fv
        node.output_shape = fv.shape
        # In most cases, the batch dimension is just the first dimension
        # if the node depends on input. Otherwise if the node doesn't
        # depend on input, there is no batch dimension (default is -1).
        node.batch_dim = 0 if node.from_input else node.batch_dim
        # Unperturbed node but it is not a root node.
        # Save forward_value to value. (Can be used in forward bounds.)
        if not node.from_input and len(node.inputs) > 0:
            node.value = node.forward_value
        return fv

    def forward(self, *x, final_node_name=None,
                interm_bounds=None,
                clear_forward_only=False,
                reset_perturbed_nodes=True,
                cache_bounds=False):
        r"""Standard forward computation for the network.

        Args:
            x (tuple or None): Input to the model.

            final_node_name (str, optional): The name of the final node in the
            model. The value on the corresponding node will be returned.

            clear_forward_only (bool, default `False`): Whether only standard
            forward values stored on the nodes should be cleared. If `True`,
            only standard forward values stored on the nodes will be cleared.
            Otherwise, bound information on the nodes will also be cleared.

            reset_perturbed_nodes (bool, default `True`): Mark all perturbed
            nodes with input perturbations. When set to `True`, it may
            accidentally clear all .perturbed properties for intermediate
            nodes.

        Returns:
            output: The output of the model, or if `final_node_name` is not
            `None`, return the value on the corresponding node instead.
        """
        self.set_input(*x,
                       interm_bounds=interm_bounds,
                       clear_forward_only=clear_forward_only,
                       reset_perturbed_nodes=reset_perturbed_nodes,
                       cache_bounds=cache_bounds)
        if final_node_name is None:
            final_node_name = self.output_name[0]
        return self.get_forward_value(self[final_node_name])

    def _mark_perturbed_nodes(self, input):
        """Mark the graph nodes and determine which nodes need perturbation."""
        # Set some of the input as perturbed if they are bounded objects
        any_perturbed = False
        for name, index in zip(self.input_name, self.input_index):
            if index is None:
                continue
            if isinstance(input[index], (BoundedTensor, BoundedParameter)):
                self[name].perturbed = True
                any_perturbed = True
        # If none of the inputs is a bounded object, set all of them as perturbed
        if not any_perturbed:
            for name, index in zip(self.input_name, self.input_index):
                if index is not None:
                    self[name].perturbed = True

        degree_in = {}
        queue = deque()
        relus = []
        # Initially the queue contains all "root" nodes.
        for key in self._modules.keys():
            l = self[key]
            degree_in[l.name] = len(l.inputs)
            if degree_in[l.name] == 0:
                queue.append(l)  # in_degree ==0 -> root node

        while len(queue) > 0:
            node = queue.popleft()
            # We set the relu here to ensure the list is sorted according to topological order.
            if isinstance(node, BoundRelu):
                relus.append(node)
            # Obtain all output node, and add the output nodes to the queue if
            # all its input nodes have been visited.
            # The initial "perturbed" property is set in BoundInput or
            # BoundParams object, depending on ptb.
            for name_next in node.output_name:
                node_next = self[name_next]
                if not node_next.never_perturbed:
                    # The next node is perturbed if it is already perturbed,
                    # or this node is perturbed.
                    node_next.perturbed = node_next.perturbed or node.perturbed
                degree_in[name_next] -= 1
                # all inputs of this node have been visited,
                # now put it in queue.
                if degree_in[name_next] == 0:
                    queue.append(node_next)
            node.update_requires_input_bounds()

        self.relus = relus
        self.get_optimizable_activations()
        self.splittable_activations = self.get_splittable_activations()
        self.perturbed_optimizable_activations = (
            self.get_perturbed_optimizable_activations())
        return

    def _check_patches_mode(self):
        """Disable patches mode if there is no Conv node.

        This is a workaround (before a more general patches mode is implemented)
        to avoid issues relevant to the patches node,
        for complicated models without any Conv.
        """
        has_conv = False
        for node in self.nodes():
            if isinstance(node, (BoundConv, BoundConvTranspose, BoundConv2dGrad)):
                has_conv = True
        if not has_conv and self.conv_mode == 'patches':
            self.conv_mode = 'matrix'
            for node in self.nodes():
                if getattr(node, 'mode', None) == 'patches':
                    node.mode = 'matrix'

    def _clear_and_set_new(
        self,
        interm_bounds,
        clear_forward_only=False,
        reset_perturbed_nodes=True,
        cache_bounds=False,
    ):
        for l in self.nodes():
            if hasattr(l, 'linear'):
                if isinstance(l.linear, tuple):
                    for item in l.linear:
                        del item
                delattr(l, 'linear')

            if hasattr(l, 'patch_size'):
                l.patch_size = {}

            if clear_forward_only:
                if hasattr(l, 'forward_value'):
                    delattr(l, 'forward_value')
            else:
                for attr in ['interval', 'forward_value', 'd',
                             'lA', 'lower_d', 'upper_k']:
                    if hasattr(l, attr):
                        delattr(l, attr)
                if cache_bounds:
                    l.move_lower_and_upper_bounds_to_cache()
                else:
                    l.delete_lower_and_upper_bounds()

            for attr in ['zero_backward_coeffs_l', 'zero_backward_coeffs_u',
                         'zero_lA_mtx', 'zero_uA_mtx']:
                setattr(l, attr, False)
            # Given an interval here to make IBP/CROWN start from this node
            if interm_bounds is not None and l.name in interm_bounds.keys():
                l.interval = tuple(interm_bounds[l.name][:2])
                l.lower = interm_bounds[l.name][0]
                l.upper = interm_bounds[l.name][1]
                if l.lower is not None:
                    l.lower = l.lower.detach().requires_grad_(False)
                if l.upper is not None:
                    l.upper = l.upper.detach().requires_grad_(False)
            # Mark all nodes as non-perturbed except for weights.
            if reset_perturbed_nodes:
                if not hasattr(l, 'perturbation') or l.perturbation is None:
                    l.perturbed = False

            # Clear operator-specific attributes
            l.clear()

    def set_input(
        self,
        *x,
        interm_bounds=None,
        clear_forward_only=False,
        reset_perturbed_nodes=True,
        cache_bounds=False,
    ):
        self._clear_and_set_new(
            interm_bounds=interm_bounds,
            clear_forward_only=clear_forward_only,
            reset_perturbed_nodes=reset_perturbed_nodes,
            cache_bounds=cache_bounds,
        )
        inputs_unpacked = unpack_inputs(x)
        for name, index in zip(self.input_name, self.input_index):
            if index is None:
                continue
            node = self[name]
            node.value = inputs_unpacked[index]
            if isinstance(node.value, (BoundedTensor, BoundedParameter)):
                node.perturbation = node.value.ptb
            else:
                node.perturbation = None
        # Mark all perturbed nodes.
        if reset_perturbed_nodes:
            self._mark_perturbed_nodes(inputs_unpacked)

    def _get_node_input(self, nodesOP, nodesIn, node):
        ret = []
        for i in range(len(node.inputs)):
            for op in nodesOP:
                if op.name == node.inputs[i]:
                    ret.append(op.bound_node)
                    break
            if len(ret) == i + 1:
                continue
            for io in nodesIn:
                if io.name == node.inputs[i]:
                    ret.append(io.bound_node)
                    break
            if len(ret) <= i:
                raise ValueError(f'cannot find inputs of node: {node.name}')
        return ret

    def _to(self, obj, dest, inplace=False):
        """ Move all tensors in the object to a specified dest
        (device or dtype). The inplace=True option is available for dict."""
        if obj is None:
            return obj
        elif isinstance(obj, torch.Tensor):
            return obj.to(dest)
        elif isinstance(obj, Patches):
            return obj.patches.to(dest)
        elif isinstance(obj, tuple):
            return tuple([self._to(item, dest) for item in obj])
        elif isinstance(obj, list):
            return list([self._to(item, dest) for item in obj])
        elif isinstance(obj, dict):
            if inplace:
                for k, v in obj.items():
                    obj[k] = self._to(v, dest, inplace=True)
                return obj
            else:
                return {k: self._to(v, dest) for k, v in obj.items()}
        else:
            raise NotImplementedError(type(obj))

    def _convert_nodes(self, model, global_input):
        r"""
        Returns:
            nodesOP (list): List of operator nodes
            nodesIn (list): List of input nodes
            nodesOut (list): List of output nodes
            template (object): Template to specify the output format
        """
        global_input_cpu = self._to(global_input, 'cpu')
        if self.ori_training:
            model.train()
        else:
            model.eval()
        model.to('cpu')
        nodesOP, nodesIn, nodesOut, template = parse_module(
            model, global_input_cpu)
        model.to(self.device)
        for i in range(0, len(nodesIn)):
            if nodesIn[i].param is not None:
                nodesIn[i] = nodesIn[i]._replace(
                    param=nodesIn[i].param.to(self.device))

        # Convert input nodes and parameters.
        attr = {'device': self.device}
        for i, n in enumerate(nodesIn):
            if n.input_index is not None:
                nodesIn[i] = nodesIn[i]._replace(bound_node=BoundInput(
                    ori_name=nodesIn[i].ori_name,
                    value=global_input[nodesIn[i].input_index],
                    perturbation=nodesIn[i].perturbation,
                    input_index=n.input_index, options=self.bound_opts,
                    attr=attr))
            else:
                bound_class = BoundParams if isinstance(
                    nodesIn[i].param, nn.Parameter) else BoundBuffers
                nodesIn[i] = nodesIn[i]._replace(bound_node=bound_class(
                    ori_name=nodesIn[i].ori_name, value=nodesIn[i].param,
                    perturbation=nodesIn[i].perturbation, options=self.bound_opts,
                    attr=attr))

        unsupported_ops = []

        # Convert other operation nodes.
        for n in range(len(nodesOP)):
            attr = nodesOP[n].attr
            inputs = self._get_node_input(nodesOP, nodesIn, nodesOP[n])
            try:
                if nodesOP[n].op in self.custom_ops:
                    op = self.custom_ops[nodesOP[n].op]
                elif nodesOP[n].op in bound_op_map:
                    op = bound_op_map[nodesOP[n].op]
                elif nodesOP[n].op.startswith('aten::ATen'):
                    op = globals()[f'BoundATen{attr["operator"].capitalize()}']
                elif nodesOP[n].op.startswith('onnx::'):
                    op = globals()[f'Bound{nodesOP[n].op[6:]}']
                else:
                    raise KeyError
            except (NameError, KeyError):
                unsupported_ops.append(nodesOP[n])
                logger.error('The node has an unsupported operation: %s',
                             nodesOP[n])
                continue
            attr['device'] = self.device

            # FIXME generalize
            if (nodesOP[n].op == 'onnx::BatchNormalization'
                    or getattr(op, 'TRAINING_FLAG', False)):
                # BatchNormalization node needs model.training flag to set
                # running mean and vars set training=False to avoid wrongly
                # updating running mean/vars during bound wrapper
                nodesOP[n] = nodesOP[n]._replace(bound_node=op(
                    attr, inputs, nodesOP[n].output_index, self.bound_opts,
                    False))
            else:
                nodesOP[n] = nodesOP[n]._replace(bound_node=op(
                    attr, inputs, nodesOP[n].output_index, self.bound_opts))

        if unsupported_ops:
            logger.error('Unsupported operations:')
            for n in unsupported_ops:
                logger.error(f'Name: {n.op}, Attr: {n.attr}')
            raise NotImplementedError('There are unsupported operations')

        for node in nodesIn + nodesOP:
            node.bound_node.name = node.name

        nodes_dict = {}
        for node in nodesOP + nodesIn:
            nodes_dict[node.name] = node.bound_node
        nodesOP = [n.bound_node for n in nodesOP]
        nodesIn = [n.bound_node for n in nodesIn]
        nodesOut = [nodes_dict[n] for n in nodesOut]

        return nodesOP, nodesIn, nodesOut, template

    def _build_graph(self, nodesOP, nodesIn, nodesOut, template):
        # We were assuming that the original model had only one output node.
        assert len(nodesOut) == 1
        self.final_name = nodesOut[0].name
        self.input_name, self.input_index, self.root_names = [], [], []
        self.output_name = [n.name for n in nodesOut]
        self.output_template = template
        self._modules.clear()
        for node in nodesIn:
            self.add_input_node(node, index=node.input_index)
        self.add_nodes(nodesOP)
        if self.conv_mode == 'patches':
            self.root_names: List[str] = [node.name for node in nodesIn]

    def rename_nodes(self, nodesOP, nodesIn, rename_dict):
        def rename(node):
            node.name = rename_dict[node.name]
            return node
        for i in range(len(nodesOP)):
            nodesOP[i] = rename(nodesOP[i])
        for i in range(len(nodesIn)):
            nodesIn[i] = rename(nodesIn[i])

    def _split_complex(self, nodesOP, nodesIn):
        finished = True
        for n in range(len(nodesOP)):
            if hasattr(nodesOP[n], 'complex') and nodesOP[n].complex:
                complex_node = nodesOP[n]

                finished = False
                _nodesOP, _nodesIn, _nodesOut, _ = self._convert_nodes(
                    nodesOP[n].model, nodesOP[n].input)
                # assuming each supported complex operation only has one output
                assert len(_nodesOut) == 1

                name_base = nodesOP[n].name + '/split'
                rename_dict = {}
                for node in _nodesOP + _nodesIn:
                    rename_dict[node.name] = name_base + node.name
                num_inputs = len(nodesOP[n].inputs)
                for i in range(num_inputs):
                    rename_dict[_nodesIn[i].name] = nodesOP[n].input_name[i]
                rename_dict[_nodesOP[-1].name] = nodesOP[n].name

                self.rename_nodes(_nodesOP, _nodesIn, rename_dict)

                output_name = _nodesOP[-1].name
                # Any input node of some node within the complex node should be
                # replaced with the complex node's corresponding input node.
                for node in _nodesOP:
                    for i in range(len(node.inputs)):
                        if node.input_name[i] in nodesOP[n].input_name:
                            index = nodesOP[n].input_name.index(
                                node.input_name[i])
                            node.inputs[i] = nodesOP[n].inputs[index]
                # For any output node of this complex node,
                # modify its input node.
                for node in nodesOP:
                    if output_name in node.input_name:
                        index = node.input_name.index(output_name)
                        node.inputs[index] = _nodesOP[-1]
                # Mark where the nodes come from
                for node in _nodesOP:
                    node.from_complex_node = type(complex_node).__name__

                nodesOP = nodesOP[:n] + _nodesOP + nodesOP[(n + 1):]
                nodesIn = nodesIn + _nodesIn[num_inputs:]

                break

        return nodesOP, nodesIn, finished

    def _get_node_name_map(self):
        """Build a dict with {ori_name: name, name: ori_name}"""
        self.node_name_map = {}
        for node in self.nodes():
            if isinstance(node, (BoundInput, BoundParams)):
                for p in list(node.named_parameters()):
                    if node.ori_name not in self.node_name_map:
                        name = f'{node.name}.{p[0]}'
                        self.node_name_map[node.ori_name] = name
                        self.node_name_map[name] = node.ori_name
                for p in list(node.named_buffers()):
                    if node.ori_name not in self.node_name_map:
                        name = f'{node.name}.{p[0]}'
                        self.node_name_map[node.ori_name] = name
                        self.node_name_map[name] = node.ori_name

    # convert a Pytorch model to a model with bounds
    def _convert(self, model, global_input):
        if self.verbose:
            logger.info('Converting the model...')

        self.num_global_inputs = len(global_input)

        nodesOP, nodesIn, nodesOut, template = self._convert_nodes(
            model, global_input)
        global_input = self._to(global_input, self.device)

        while True:
            self._build_graph(nodesOP, nodesIn, nodesOut, template)
            self.forward(*global_input)  # running means/vars changed
            nodesOP, nodesIn, finished = self._split_complex(nodesOP, nodesIn)
            if finished:
                break

        self._get_node_name_map()

        ori_state_dict_mapped = OrderedDict()
        for k, v in self.ori_state_dict.items():
            if k in self.node_name_map:
                ori_state_dict_mapped[self.node_name_map[k]] = v
        self.load_state_dict(ori_state_dict_mapped)
        if self.ori_training:
            model.load_state_dict(self.ori_state_dict)
        delattr(self, 'ori_state_dict')

        # The name of the final node used in the last call to `compute_bounds`
        self.last_final_node_name = None

        if self.verbose:
            logger.info('Model converted to support bounds')

    def check_prior_bounds(self, node, C=None):
        if node.prior_checked or not (node.used and node.perturbed):
            return
        if C is not None and isinstance(node, BoundConcat):
            # If the last node is a BoundConcat, it's possible that only some of
            # the input nodes of the BoundConcat are needed in the specification.
            # In this case, we only check the bounds of the input nodes that are
            # actually used in the specification. All other branches are
            # considered as not used, and their bounds are not checked.
            # FIXME: In this case, node.used of some nodes may be incorrect.
            offset = 0
            assert isinstance(C, torch.Tensor) and C.ndim == 3
            C = C.abs().sum(dim=[0, 1])
            for node_input in node.inputs:
                size = prod(node_input.output_shape[1:])
                C_s = C[offset:offset+size].sum()
                if (C_s != 0).any():
                    self.check_prior_bounds(node_input)
                offset += size
        else:
            for n in node.inputs:
                self.check_prior_bounds(n)
        tighten_input_bounds = (
            self.bound_opts['optimize_bound_args']['tighten_input_bounds']
        )
        directly_optimize_layer_names = (
            self.bound_opts['optimize_bound_args']['directly_optimize']
        )
        bound_every_node = (
            self.bound_opts['bound_every_node']
        )
        for i in range(len(node.inputs)):
            if (
                i in node.requires_input_bounds
                or not node.inputs[i].perturbed
                or node.inputs[i].name in self.layers_with_constraint
                # allows to tighten input bounds
                or (isinstance(node.inputs[i], BoundInput) and tighten_input_bounds)
                # layers whos optimization is forced
                # (for consecutive layers introduced as part of invprop)
                or node.inputs[i].name in directly_optimize_layer_names
                or bound_every_node
            ):
                self.compute_intermediate_bounds(
                    node.inputs[i], prior_checked=True)
        node.prior_checked = True

    def compute_intermediate_bounds(self, node: Bound, prior_checked=False):
        tighten_input_bounds = (
            self.bound_opts['optimize_bound_args']['tighten_input_bounds']
        )
        directly_optimize_layer_names = (
            self.bound_opts['optimize_bound_args']['directly_optimize']
        )
        best_of_oc_and_no_oc = (
            self.bound_opts['optimize_bound_args']['best_of_oc_and_no_oc']
        )
        if (
            node.is_lower_bound_current()
            and not (
                isinstance(node, BoundInput) and tighten_input_bounds
                or node.name in directly_optimize_layer_names
            )
        ):
            if node.name in self.layers_with_constraint:
                node.clamp_interim_bounds()
            return

        logger.debug(f'Getting the bounds of {node}')

        if not prior_checked:
            self.check_prior_bounds(node)

        if not node.perturbed:
            fv = self.get_forward_value(node)
            node.interval = node.lower, node.upper = fv, fv
            return

        # FIXME check that weight perturbation is not affected
        #      (from_input=True should be set for weights)
        if not node.from_input and hasattr(node, 'forward_value'):
            node.lower = node.upper = self.get_forward_value(node)
            return

        reference_bounds = self.reference_bounds

        if self.use_forward:
            # forward
            node.lower, node.upper = self.forward_general(
                node=node, concretize=True)
        else:
            # backward
            if self.check_IBP_intermediate(node):
                # Intermediate bounds for some operators are directly
                # computed from their input nodes by IBP
                # (such as BoundRelu, BoundNeg)
                logger.debug('IBP propagation for intermediate bounds on %s', node)
            # For the first linear layer, IBP can give the same tightness as CROWN.
            elif not self.check_IBP_first_linear(node):
                ref_intermediate = self.get_ref_intermediate_bounds(node)
                sparse_C = self.get_sparse_C(node, ref_intermediate)
                newC, reduced_dim, unstable_idx, unstable_size = sparse_C

                # Special case for BoundRelu when sparse intermediate bounds are disabled
                # Currently sparse intermediate bounds are restricted to ReLU models only
                skip = False
                if unstable_idx is None:
                    if (len(node.output_name) == 1
                            and isinstance(self[node.output_name[0]], BoundTwoPieceLinear)
                            and node.name in self.reference_bounds):
                        lower, upper = self.reference_bounds[node.name]
                        fully_stable = torch.logical_or(lower>=0, upper<=0).all()
                        if fully_stable:
                            node.lower, node.upper = lower, upper
                            skip = True
                elif unstable_size == 0:
                    skip = True

                if not skip:
                    apply_output_constraints_to = self.bound_opts[
                        'optimize_bound_args']['apply_output_constraints_to']
                    if self.return_A:
                        node.lower, node.upper, _ = self.backward_general(
                            node, newC, unstable_idx=unstable_idx,
                            apply_output_constraints_to=apply_output_constraints_to)
                    else:
                        # Compute backward bounds only when there are unstable
                        # neurons, or when we don't know which neurons are unstable.
                        node.lower, node.upper = self.backward_general(
                            node, newC, unstable_idx=unstable_idx,
                            apply_output_constraints_to=apply_output_constraints_to)
                    if torch.any((node.upper - node.lower).abs() > 1e10):
                        if len(apply_output_constraints_to) > 0 and not best_of_oc_and_no_oc:
                            warnings.warn('Very weak bounds detected. This can potentially be '
                                'fixed by setting best_of_oc_and_no_oc=True.')

                if reduced_dim:
                    self.restore_sparse_bounds(
                        node, unstable_idx, unstable_size, ref_intermediate)

                if self.bound_opts['compare_crown_with_ibp']:
                    node.lower, node.upper = self.compare_with_IBP(node, node.lower, node.upper)

        # node.lower and node.upper (intermediate bounds) are computed in
        # the above function. If we have bound references, we set them here
        # to always obtain a better set of bounds.
        if node.name in reference_bounds:
            ref_bounds = reference_bounds[node.name]
            # Initially, the reference bound and the computed bound can be
            # exactly the same when intermediate layer beta is 0. This will
            # prevent gradients flow. So we need a small guard here.
            # Set the intermediate layer bounds using reference bounds,
            # always choosing the tighter one.
            # Assert no NaNs in reference bounds before using them
            assert not torch.isnan(ref_bounds[0]).any(), (
                f'NaN detected in reference lower bound of layer {node.name}')
            node.lower = (torch.max(ref_bounds[0], node.lower).detach()
                          - node.lower.detach() + node.lower)
            assert not torch.isnan(ref_bounds[1]).any(), (
                f'NaN detected in reference upper bound of layer {node.name}')
            node.upper = (node.upper - (node.upper.detach()
                          - torch.min(ref_bounds[1], node.upper).detach()))
            # Also update bounds in node.linear (if exist)
            if hasattr(node, 'linear'):
                node.linear.lower = node.lower
                node.linear.upper = node.upper
            # Otherwise, we only use reference bounds to check which neurons
            # are unstable.

        # prior constraint bounds
        if node.name in self.layers_with_constraint:
            node.clamp_interim_bounds()
        # FIXME (12/28): we should be consistent, and only use
        # node.interval, do not use node.lower or node.upper!
        node.interval = (node.lower, node.upper)

    def get_ref_intermediate_bounds(self, node):
        sparse_intermediate_bounds_with_ibp = self.bound_opts.get(
            'sparse_intermediate_bounds_with_ibp', True)
        # Sparse intermediate bounds can be enabled
        # if aux_reference_bounds are given.
        # (this is enabled for ReLU only, and not for other activations.)
        sparse_intermediate_bounds = (self.bound_opts.get(
            'sparse_intermediate_bounds', False)
            and isinstance(self[node.output_name[0]], BoundRelu))

        ref_intermediate_lb, ref_intermediate_ub = None, None
        if sparse_intermediate_bounds:
            if node.name not in self.aux_reference_bounds:
                # If aux_reference_bounds are not available,
                # we can use IBP to compute these bounds.
                if sparse_intermediate_bounds_with_ibp:
                    with torch.no_grad():
                        # Get IBP bounds for this layer;
                        # we set delete_bounds_after_use=True which does
                        # not save extra intermediate bound tensors.
                        ret_ibp = self.IBP_general(
                            node=node, delete_bounds_after_use=True)
                        ref_intermediate_lb = ret_ibp[0]
                        ref_intermediate_ub = ret_ibp[1]
                else:
                    sparse_intermediate_bounds = False
            else:
                aux_bounds = self.aux_reference_bounds[node.name]
                ref_intermediate_lb, ref_intermediate_ub = aux_bounds

        return sparse_intermediate_bounds, ref_intermediate_lb, ref_intermediate_ub

    def merge_A_dict(self, lA_dict, uA_dict):
        merged_A = {}
        for output_node_name in lA_dict:
            merged_A[output_node_name] = {}
            lA_dict_ = lA_dict[output_node_name]
            uA_dict_ = uA_dict[output_node_name]
            for input_node_name in lA_dict_:
                merged_A[output_node_name][input_node_name] = {
                    'lA': lA_dict_[input_node_name]['lA'],
                    'uA': uA_dict_[input_node_name]['uA'],
                    'lbias': lA_dict_[input_node_name]['lbias'],
                    'ubias': uA_dict_[input_node_name]['ubias'],
                }
        return merged_A

    def compute_bounds(
            self, x=None, aux=None, C=None, method='backward', IBP=False,
            forward=False, bound_lower=True, bound_upper=True, reuse_ibp=False,
            reuse_alpha=False, return_A=False, needed_A_dict=None,
            final_node_name=None, average_A=False,
            interm_bounds=None, reference_bounds=None,
            intermediate_constr=None, alpha_idx=None,
            aux_reference_bounds=None, need_A_only=False,
            cutter=None, decision_thresh=None,
            update_mask=None, ibp_nodes=None, cache_bounds=False):
        r"""Main function for computing bounds.

        Args:
            x (tuple or None): Input to the model. If it is None, the input
            from the last `forward` or `compute_bounds` call is reused.
            Otherwise: the number of elements in the tuple should be
            equal to the number of input nodes in the model, and each element in
            the tuple corresponds to the value for each input node respectively.
            It should look similar as the `global_input` argument when used for
            creating a `BoundedModule`.

            aux (object, optional): Auxliary information that can be passed to
            `Perturbation` classes for initializing and concretizing bounds,
            e.g., additional information for supporting synonym word subsitution
            perturbaiton.

            C (Tensor): The specification matrix that can map the output of the
            model with an additional linear layer. This is usually used for
            maping the logits output of the model to classification margins.

            method (str): The main method for bound computation. Choices:
                * `IBP`: purely use Interval Bound Propagation (IBP) bounds.
                * `CROWN-IBP`: use IBP to compute intermediate bounds,
                but use CROWN (backward mode LiRPA) to compute the bounds of the
                final node.
                * `CROWN`: purely use CROWN to compute bounds for intermediate
                nodes and the final node.
                * `Forward`: purely use forward mode LiRPA.
                * `Forward+Backward`: use forward mode LiRPA for intermediate
                nodes, but further use CROWN for the final node.
                * `CROWN-Optimized` or `alpha-CROWN`: use CROWN, and also
                optimize the linear relaxation parameters for activations.
                * `forward-optimized`: use forward bounds with optimized linear
                relaxation.
                * `dynamic-forward`: use dynamic forward bound propagation where
                new input variables may be dynamically introduced for
                nonlinearities.
                * `dynamic-forward+backward`: use dynamic forward mode for
                intermediate nodes, but use CROWN for the final node.

            IBP (bool, optional): If `True`, use IBP to compute the bounds of
            intermediate nodes. It can be automatically set according to
            `method`.

            forward (bool, optional): If `True`, use the forward mode bound
            propagation to compute the bounds of intermediate nodes. It can be
            automatically set according to `method`.

            bound_lower (bool, default `True`): If `True`, the lower bounds of
            the output needs to be computed.

            bound_upper (bool, default `True`): If `True`, the upper bounds of
            the output needs to be computed.

            reuse_ibp (bool, optional): If `True` and `method` is None, reuse
            the previously saved IBP bounds.

            final_node_name (str, optional): Set the final node in the
            computational graph for bound computation. By default, the final
            node of the originally built computational graph is used.

            return_A (bool, optional): If `True`, return linear coefficients
            in bound propagation (`A` tensors) with `needed_A_dict` set.

            needed_A_dict (dict, optional): A dictionary specifying linear
            coefficients (`A` tensors) that are needed and should be returned.
            Each key in the dictionary is the name of a starting node in
            backward bound propagation, with a list as the value for the key,
            which specifies the names of the ending nodes in backward bound
            propagation, and the linear coefficients of the starting node w.r.t.
            the specified ending nodes are returned. By default, it is empty.

            reuse_alpha (bool, optional): If `True`, reuse previously saved
            alpha values when they are not being optimized.

            decision_thresh (float, optional): In CROWN-optimized mode, we will
            use this decision_thresh to dynamically optimize those domains that
            <= the threshold.

            interm_bounds: A dictionary of 2-element tuple/list
            containing lower and upper bounds for intermediate layers.
            The dictionary keys should include the names of the layers whose
            bounds should be set without recomputation. The layer names can be
            viewed by setting environment variable AUTOLIRPA_DEBUG=1.
            The values of each dictionary elements are (lower_bounds,
            upper_bounds) where "lower_bounds" and "upper_bounds" are two
            tensors with the same shape as the output shape of this layer. If
            you only need to set intermediate layer bounds for certain layers,
            then just include these layers' names in the dictionary.

            reference_bounds: Format is similar to "interm_bounds".
            However, these bounds are only used as a reference, and the bounds
            for intermediate layers will still be computed (e.g., using CROWN,
            IBP or other specified methods). The computed bounds will be
            compared to "reference_bounds" and the tighter one between the two
            will be used.

            aux_reference_bounds: Format is similar to intermediate layer
            bounds. However, these bounds are only used for determine which
            neurons are stable and which neurons are unstable for ReLU networks.
            Unstable neurons' intermediate layer bounds will be recomputed.

            cache_bounds: If `True`, the currently set lower and upper bounds will not
            be deleted, but cached for use by the INVPROP algorithm. This should not be
            set by the user, but only in `_get_optimized_bounds`.

        Returns:
            bound (tuple): When `return_A` is `False`, return a tuple of
            the computed lower bound and upper bound. When `return_A`
            is `True`, return a tuple of lower bound, upper bound, and
            `A` dictionary.
        """
        # This method only prepares everything by setting all required parameters.
        # The main logic is located in `_compute_bounds_main`. It may be called
        # repeatedly for CROWN optimizations.
        logger.debug(f'Compute bounds with {method}')
        if needed_A_dict is None: needed_A_dict = {}
        if not bound_lower and not bound_upper:
            raise ValueError(
                'At least one of bound_lower and bound_upper must be True')

        # Several shortcuts.
        compute_optimized = False
        method = method.lower() if method is not None else method
        if method == 'ibp':
            # Pure IBP bounds.
            method, IBP = None, True
        elif method in ['ibp+backward', 'ibp+crown', 'crown-ibp']:
            method, IBP = 'backward', True
        elif method == 'crown':
            method = 'backward'
        elif method == 'forward':
            forward = True
            self.dynamic = False
        elif method == 'dynamic-forward':
            forward = True
            self.dynamic = True
        elif method == 'forward+backward' or method == 'forward+crown':
            method, forward = 'backward', True
        elif method == 'dynamic-forward+backward' or method == 'dynamic-forward+crown':
            self.dynamic = True
            method, forward = 'backward', True
        elif method in ['crown-optimized', 'alpha-crown', 'forward-optimized']:
            # Lower and upper bounds need two separate rounds of optimization.
            if method == 'forward-optimized':
                method = 'forward'
            else:
                method = 'backward'
            compute_optimized = True

        if reference_bounds is None:
            reference_bounds = {}
        if aux_reference_bounds is None:
            aux_reference_bounds = {}

        # If y in self.backward_node_pairs[x], then node y is visited when
        # doing backward bound propagation starting from node x.
        self.backward_from = dict([(node, []) for node in self._modules])

        if not bound_lower and not bound_upper:
            raise ValueError(
                'At least one of bound_lower and bound_upper in compute_bounds '
                'should be True')
        A_dict = {} if return_A else None

        if x is not None:
            if isinstance(x, torch.Tensor):
                x = (x,)
            if self.bound_opts['forward_before_compute_bounds']:
                self.forward(*x, interm_bounds=interm_bounds, cache_bounds=cache_bounds)
            else:
                self.set_input(*x, interm_bounds=interm_bounds, cache_bounds=cache_bounds)

        roots = self.roots()
        batch_size = roots[0].value.shape[0]
        dim_in = 0

        for i in range(len(roots)):
            value = roots[i].forward()
            if getattr(roots[i], 'perturbation', None) is not None:
                ret_init = roots[i].perturbation.init(
                    value, aux=aux, forward=forward)
                roots[i].linear, roots[i].center, roots[i].aux = ret_init
                # This input/parameter has perturbation.
                # Create an interval object.
                roots[i].interval = Interval(
                    roots[i].linear.lower, roots[i].linear.upper,
                    ptb=roots[i].perturbation)
                if forward:
                    roots[i].dim = roots[i].linear.lw.shape[1]
                    dim_in += roots[i].dim

            else:
                # This input/parameter does not has perturbation.
                # Use plain tuple defaulting to Linf perturbation.
                roots[i].interval = (value, value)
                roots[i].forward_value = roots[i].value = value
                roots[i].center = roots[i].lower = roots[i].upper = value

            roots[i].lower, roots[i].upper = roots[i].interval

        if forward:
            self.init_forward(roots, dim_in)

        for n in self.nodes():
            if isinstance(n, BoundRelu):
                for node in n.inputs:
                    if hasattr(node, 'relu_followed'):
                        node.relu_followed = True

            # Inject update mask inside the activations
            # update_mask: None or bool tensor([batch_size])
            # If set to a tensor, only update the alpha and beta of selected
            # element (with element=1).
            n.alpha_beta_update_mask = update_mask

        final = (self.final_node() if final_node_name is None
                 else self[final_node_name])
        # BFS to find out whether each node is used given the current final node
        self._set_used_nodes(final)

        self.use_forward = forward
        self.batch_size = batch_size
        self.dim_in = dim_in
        self.return_A = return_A
        self.A_dict = A_dict
        self.needed_A_dict = needed_A_dict
        self.intermediate_constr = intermediate_constr
        self.reference_bounds = reference_bounds
        self.aux_reference_bounds = aux_reference_bounds
        self.final_node_name = final.name
        self.ibp_nodes = ibp_nodes

        if compute_optimized:
            kwargs = dict(x=x, C=C, method=method, interm_bounds=interm_bounds,
                reference_bounds=reference_bounds, return_A=return_A,
                aux_reference_bounds=aux_reference_bounds,
                needed_A_dict=needed_A_dict,
                final_node_name=final_node_name,
                cutter=cutter, decision_thresh=decision_thresh)
            if bound_upper:
                ret2 = self._get_optimized_bounds(bound_side='upper', **kwargs)
            else:
                ret2 = None
            if bound_lower:
                ret1 = self._get_optimized_bounds(bound_side='lower', **kwargs)
            else:
                ret1 = None
            if bound_lower and bound_upper:
                if return_A:
                    # Needs to merge the A dictionary.
                    return ret1[0], ret2[1], self.merge_A_dict(ret1[2], ret2[2])
                else:
                    return ret1[0], ret2[1]
            elif bound_lower:
                return ret1  # ret1[1] is None.
            elif bound_upper:
                return ret2  # ret2[0] is None.

        return self._compute_bounds_main(C=C,
                                         method=method,
                                         IBP=IBP,
                                         bound_lower=bound_lower,
                                         bound_upper=bound_upper,
                                         reuse_ibp=reuse_ibp,
                                         reuse_alpha=reuse_alpha,
                                         average_A=average_A,
                                         alpha_idx=alpha_idx,
                                         need_A_only=need_A_only,
                                         update_mask=update_mask)

    def save_intermediate(self, save_path=None):
        r"""A function for saving intermediate bounds.

        Please call this function after `compute_bounds`, or it will output
        IBP bounds by default.

        Args:
            save_path (str, default `None`): If `None`, the intermediate bounds
            will not be saved, or it will be saved at the designated path.

        Returns:
            save_dict (dict): Return a dictionary of lower and upper bounds, with
            the key being the name of the layer.
        """
        save_dict = OrderedDict()
        for node in self.nodes():
            if node.used and node.perturbed:
                if not hasattr(node, 'interval'):
                    ibp_lower, ibp_upper = self.IBP_general(node,
                        delete_bounds_after_use=True)
                    dim_output = int(prod(node.output_shape[1:]))
                    C = torch.eye(dim_output, device=self.device).expand(
                        self.batch_size, dim_output, dim_output)
                    crown_lower, crown_upper = self.backward_general(node, C=C)
                    save_dict[node.name] = (
                        torch.max(crown_lower, ibp_lower),
                        torch.min(crown_upper, ibp_upper))
                else:
                    save_dict[node.name] = (node.lower, node.upper)

        if save_path is not None:
            torch.save(save_dict, save_path)
        return save_dict

    def _compute_bounds_main(self, C=None, method='backward', IBP=False,
            bound_lower=True, bound_upper=True, reuse_ibp=False,
            reuse_alpha=False, average_A=False, alpha_idx=None,
            need_A_only=False, update_mask=None):
        """The core implementation of compute_bounds.

        Seperated because compute_bounds may call _get_optimized_bounds which
        repeatedly calls this method. Otherwise, the preprocessing done in
        compute_bounds would be executed for each iteration.
        """

        final = (self.final_node() if self.final_node_name is None
                 else self[self.final_node_name])
        logger.debug(f'Final node {final.__class__.__name__}({final.name})')

        if IBP and method is None and reuse_ibp:
            # directly return the previously saved ibp bounds
            return self.ibp_lower, self.ibp_upper

        if IBP:
            self.ibp_lower, self.ibp_upper = self.IBP_general(node=final, C=C)

        if method is None:
            return self.ibp_lower, self.ibp_upper

        # TODO: if compute_bounds is called with a method that causes alphas to be
        # optimized, C will be allocated in each iteration. We could allocate it once
        # in compute_bounds, but e.g. `IBP_general` and code in `_get_optimized_bounds`
        # relies on the fact that it can be None
        if C is None:
            # C is an identity matrix by default
            if final.output_shape is None:
                raise ValueError(
                    f'C is not missing while node {final} has no default shape')
            dim_output = int(prod(final.output_shape[1:]))
            # TODO: use an eyeC object here.
            C = torch.eye(dim_output, device=self.device).expand(
                self.batch_size, dim_output, dim_output)

        # Reuse previously saved alpha values,
        # even if they are not optimized now
        # This must be done here instead of `compute_bounds`, as other code might change
        # it (e.g. `_get_optimized_bounds`)
        if reuse_alpha:
            self.opt_reuse()
        else:
            self.opt_no_reuse()

        for node in self.nodes():
            # All nodes may need to be recomputed
            node.prior_checked = False

        self.check_prior_bounds(final, C=C)

        if method == 'backward':
            apply_output_constraints_to = (
                self.bound_opts['optimize_bound_args']['apply_output_constraints_to']
            )
            # This is for the final output bound.
            # No need to pass in intermediate layer beta constraints.
            ret = self.backward_general(
                final, C,
                bound_lower=bound_lower, bound_upper=bound_upper,
                average_A=average_A, need_A_only=need_A_only,
                unstable_idx=alpha_idx, update_mask=update_mask,
                apply_output_constraints_to=apply_output_constraints_to)

            if self.bound_opts['compare_crown_with_ibp']:
                new_lower, new_upper = self.compare_with_IBP(final, lower=ret[0], upper=ret[1], C=C)
                ret = (new_lower, new_upper) + ret[2:]

            # FIXME when C is specified, lower and upper should not be saved to
            # final.lower and final.upper, because they are not the bounds for
            # the node.
            final.lower, final.upper = ret[0], ret[1]

            return ret
        elif method == 'forward' or method == 'dynamic-forward':
            return self.forward_general(C=C, node=final, concretize=True)
        else:
            raise NotImplementedError

    def _set_used_nodes(self, final):
        # By default, all *.used are initialized to False.
        # We set the used nodes by BFS from the final node.
        if final.name != self.last_final_node_name:
            self.last_final_node_name = final.name
            final.used = True
            queue = deque([final])
            while len(queue) > 0:
                n = queue.popleft()
                for n_pre in n.inputs:
                    if not n_pre.used:
                        n_pre.used = True
                        queue.append(n_pre)
        # Based on "used" and "perturbed" properties, find out which
        # layer requires intermediate layer bounds.
        self.layers_requiring_bounds = self.get_layers_requiring_bounds()

    def init_infeasible_bounds_constraints(self, batchsize, device):
        '''Simply initialize the infeasible bound record.'''
        self.infeasible_bounds_constraints = torch.full((batchsize,), False, device=device)

    from .interval_bound import (
        IBP_general, _IBP_loss_fusion, check_IBP_intermediate,
        check_IBP_first_linear, compare_with_IBP)
    from .forward_bound import (
        forward_general, forward_general_dynamic, forward_refinement, init_forward)
    from .backward_bound import (
        backward_general, get_sparse_C,
        check_optimized_variable_sparsity, restore_sparse_bounds,
        get_alpha_crown_start_nodes, get_unstable_locations, batched_backward,
        _preprocess_C)
    from .output_constraints import (
        backward_general_with_output_constraint, invprop_enabled,
        backward_general_invprop, invprop_init_infeasible_bounds,
        invprop_check_infeasible_bounds)
    from .optimized_bounds import (
        _get_optimized_bounds, init_alpha, update_best_beta,
        opt_reuse, opt_no_reuse, _to_float64, _to_default_dtype)
    from .beta_crown import (beta_crown_backward_bound, reset_beta, set_beta,
                             set_beta_cuts, get_split_nodes)
    from .jacobian import (compute_jacobian_bounds, _expand_jacobian)
    from .optimize_graph import _optimize_graph
    from .edit_graph import add_nodes, add_input_node, delete_node, replace_node
    from .tools import visualize
    from .concretize_bounds import (
        concretize_bounds, concretize_root, backward_concretize, forward_concretize)


    from .solver_module import (
        build_solver_module, _build_solver_input, _build_solver_general,
        _reset_solver_vars, _reset_solver_model)


================================================
FILE: auto_LiRPA/bound_multi_gpu.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################
from torch.nn import DataParallel
from .perturbations import *
from .bounded_tensor import BoundedTensor
from itertools import chain

class BoundDataParallel(DataParallel):
    # https://github.com/huanzhang12/CROWN-IBP/blob/master/bound_layers.py
    # This is a customized DataParallel class for our project
    def __init__(self, *inputs, **kwargs):
        super(BoundDataParallel, self).__init__(*inputs, **kwargs)
        self._replicas = None

    # Overide the forward method
    def forward(self, *inputs, **kwargs):
        disable_multi_gpu = False  # forward by single GPU
        no_replicas = False  # forward by multi GPUs but without replicate
        if "disable_multi_gpu" in kwargs:
            disable_multi_gpu = kwargs["disable_multi_gpu"]
            kwargs.pop("disable_multi_gpu")

        if "no_replicas" in kwargs:
            no_replicas = kwargs["no_replicas"]
            kwargs.pop("no_replicas")

        if not self.device_ids or disable_multi_gpu:
            if kwargs.pop("get_property", False):
                return self.get_property(self, *inputs, **kwargs)
            return self.module(*inputs, **kwargs)

        if kwargs.pop("get_property", False):
            if self._replicas is None:
                assert 0, 'please call IBP/CROWN before get_property'
            if len(self.device_ids) == 1:
                return self.get_property(self.module, **kwargs)
            inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
            kwargs = list(kwargs)
            for i in range(len(kwargs)):
                kwargs[i]['model'] = self._replicas[i]
            outputs = self.parallel_apply([self.get_property] * len(kwargs), inputs, kwargs)
            return self.gather(outputs, self.output_device)

        # Only replicate during forward/IBP propagation. Not during interval bounds
        # and CROWN-IBP bounds, since weights have not been updated. This saves 2/3
        # of communication cost.
        if not no_replicas:
            if self._replicas is None:  # first time
                self._replicas = self.replicate(self.module, self.device_ids)
            elif kwargs.get("method_opt", "forward") == "forward":
                self._replicas = self.replicate(self.module, self.device_ids)
            elif kwargs.get("x") is not None and kwargs.get("IBP") is True:  #
                self._replicas = self.replicate(self.module, self.device_ids)
            # Update the input nodes to the ones within each replica respectively
            for bounded_module in self._replicas:
                for node in bounded_module._modules.values():
                    node.inputs = [bounded_module[name] for name in node.input_name]

        for t in chain(self.module.parameters(), self.module.buffers()):
            if t.device != self.src_device_obj:
                raise RuntimeError("module must have its parameters and buffers "
                                   "on device {} (device_ids[0]) but found one of "
                                   "them on device: {}".format(self.src_device_obj, t.device))

        # TODO: can be done in parallel, only support same ptb for all inputs per forward/IBP propagation
        if len(inputs) > 0 and hasattr(inputs[0], 'ptb') and inputs[0].ptb is not None:
            # compute bounds without x
            # inputs_scatter is a normal tensor, we need to assign ptb to it if inputs is a BoundedTensor
            inputs_scatter, kwargs = self.scatter((inputs, inputs[0].ptb.x_L, inputs[0].ptb.x_U), kwargs,
                                                  self.device_ids)
            # inputs_scatter = inputs_scatter[0]
            bounded_inputs = []
            for input_s in inputs_scatter:  # GPU numbers
                # FIXME other perturbations are not supported yet
                assert isinstance(inputs[0].ptb, PerturbationLpNorm)
                ptb = PerturbationLpNorm(norm=inputs[0].ptb.norm, eps=inputs[0].ptb.eps, x_L=input_s[1], x_U=input_s[2])
                input_s = list(input_s[0])
                input_s[0] = BoundedTensor(input_s[0], ptb)
                input_s = tuple(input_s)
                bounded_inputs.append(input_s)

            # bounded_inputs = tuple(bounded_inputs)
        elif kwargs.get("x") is not None and hasattr(kwargs.get("x")[0], 'ptb') and kwargs.get("x")[0].ptb is not None:
            # compute bounds with x
            # kwargs['x'] is a normal tensor, we need to assign ptb to it
            x = kwargs.get("x")[0]
            bounded_inputs = []
            inputs_scatter, kwargs = self.scatter((inputs, x.ptb.x_L, x.ptb.x_U), kwargs, self.device_ids)
            for input_s, kw_s in zip(inputs_scatter, kwargs):  # GPU numbers
                # FIXME other perturbations are not supported yet
                assert isinstance(x.ptb, PerturbationLpNorm)
                ptb = PerturbationLpNorm(norm=x.ptb.norm, eps=x.ptb.eps, x_L=input_s[1], x_U=input_s[2])
                kw_s['x'] = list(kw_s['x'])
                kw_s['x'][0] = BoundedTensor(kw_s['x'][0], ptb)
                kw_s['x'] = (kw_s['x'])
                bounded_inputs.append(tuple(input_s[0], ))
        else:
            # normal forward
            inputs_scatter, kwargs = self.scatter(inputs, kwargs, self.device_ids)
            bounded_inputs = inputs_scatter

        if len(self.device_ids) == 1:
            return self.module(*bounded_inputs[0], **kwargs[0])
        outputs = self.parallel_apply(self._replicas[:len(bounded_inputs)], bounded_inputs, kwargs)
        return self.gather(outputs, self.output_device)

    @staticmethod
    def get_property(model, node_class=None, att_name=None, node_name=None):
        if node_name:
            # Find node by name
            # FIXME If we use `model.named_modules()`, the nodes have the
            # `BoundedModule` type rather than bound nodes.
            for node in model._modules.values():
                if node.name == node_name:
                    return getattr(node, att_name)
        else:
            # Find node by class
            for _, node in model.named_modules():
                # Find the Exp neuron in computational graph
                if isinstance(node, node_class):
                    return getattr(node, att_name)

    def state_dict(self, destination=None, prefix='', keep_vars=False):
        # add 'module.' here before each keys in self.module.state_dict() if needed
        return self.module.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)

    def _named_members(self, get_members_fn, prefix='', recurse=True, remove_duplicate: bool = True):
        return self.module._named_members(get_members_fn, prefix, recurse, remove_duplicate)

    def __getitem__(self, name):
        return self.module[name]


================================================
FILE: auto_LiRPA/bound_op_map.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################
from .bound_ops import *

bound_op_map = {
    'onnx::Gemm': BoundLinear,
    'prim::Constant': BoundPrimConstant,
    'grad::Concat': BoundConcatGrad,
    'grad::Relu': BoundReluGrad,
    'grad::Conv2d': BoundConv2dGrad,
    'grad::Slice': BoundSliceGrad,
    'grad::Sqr': BoundSqr,
    'grad::jacobian': BoundJacobianOP,
    'grad::Tanh': BoundTanhGrad,
    'grad::Sigmoid': BoundSigmoidGrad,
    'custom::Gelu': BoundGelu,
    'onnx::Clip': BoundHardTanh
}

def register_custom_op(op_name: str, bound_obj: Bound) -> None:
    bound_op_map[op_name] = bound_obj

def unregister_custom_op(op_name: str) -> None:
    bound_op_map.pop(op_name)


================================================
FILE: auto_LiRPA/bound_ops.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################
from .operators import *


================================================
FILE: auto_LiRPA/bounded_tensor.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################
import copy
import torch.nn as nn
from torch import Tensor
import torch._C as _C


class BoundedTensor(Tensor):
    @staticmethod
    # We need to override the __new__ method since Tensor is a C class
    def __new__(cls, x, ptb=None, *args, **kwargs):
        if isinstance(x, Tensor):
            tensor = super().__new__(cls, [], *args, **kwargs)
            tensor.data = x.data
            tensor.requires_grad = x.requires_grad
            return tensor
        else:
            return super().__new__(cls, x, *args, **kwargs)

    def __init__(self, x, ptb=None):
        self.ptb = ptb

    def __repr__(self):
        if hasattr(self, 'ptb') and self.ptb is not None:
            return '<BoundedTensor: {}, {}>'.format(super().__repr__(), self.ptb.__repr__())
        else:
            return '<BoundedTensor: {}, no ptb>'.format(super().__repr__())

    def clone(self, *args, **kwargs):
        tensor = BoundedTensor(super().clone(*args, **kwargs), copy.deepcopy(self.ptb))
        return tensor

    def _func(self, func, *args, **kwargs):
        temp = func(*args, **kwargs)
        new_obj = BoundedTensor([], self.ptb)
        new_obj.data = temp.data
        new_obj.requires_grad = temp.requires_grad
        return new_obj

    # Copy to other devices with perturbation
    def to(self, *args, **kwargs):
        # FIXME add a general "to" function in perturbation class, not here.
        if hasattr(self.ptb, 'x_L') and isinstance(self.ptb.x_L, Tensor):
            self.ptb.x_L = self.ptb.x_L.to(*args, **kwargs)
        if hasattr(self.ptb, 'x_U') and isinstance(self.ptb.x_U, Tensor):
            self.ptb.x_U = self.ptb.x_U.to(*args, **kwargs)
        if hasattr(self.ptb, 'eps') and isinstance(self.ptb.eps, Tensor):
            self.ptb.eps = self.ptb.eps.to(*args, **kwargs)
        return self._func(super().to, *args, **kwargs)

    @classmethod
    def _convert(cls, ret):
        if cls is Tensor:
            return ret

        if isinstance(ret, Tensor):
            if True:
                # The current implementation does not seem to need non-leaf BoundedTensor
                return ret
            else:
                # Enable this branch if non-leaf BoundedTensor should be kept
                ret = ret.as_subclass(cls)

        if isinstance(ret, tuple):
            ret = tuple(cls._convert(r) for r in ret)

        return ret

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}

        if not all(issubclass(cls, t) for t in types):
            return NotImplemented

        with _C.DisableTorchFunction():
            ret = func(*args, **kwargs)
            return cls._convert(ret)


class BoundedParameter(nn.Parameter):
    def __new__(cls, data, ptb, requires_grad=True):
        return BoundedTensor._make_subclass(cls, data, requires_grad)

    def __init__(self, data, ptb, requires_grad=True):
        self.ptb = ptb
        self.requires_grad = requires_grad

    def __deepcopy__(self, memo):
        if id(self) in memo:
            return memo[id(self)]
        else:
            result = type(self)(self.data.clone(), self.ptb, self.requires_grad)
            memo[id(self)] = result
            return result

    def __repr__(self):
        return 'BoundedParameter containing:\n{}\n{}'.format(
            self.data.__repr__(), self.ptb.__repr__())

    def __reduce_ex__(self, proto):
        raise NotImplementedError


================================================
FILE: auto_LiRPA/concretize_bounds.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################
import torch

from .utils import eyeC
from .bound_ops import *
from .patches import Patches
from .perturbations import PerturbationLpNorm

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from .bound_general import BoundedModule


def concretize_bounds(
    self: 'BoundedModule',
    node,
    lower,
    upper,
    concretize_mode='backward',
    # for `backward_concretize`
    batch_size=None,
    output_dim=None,
    average_A=None,
    # for `forward_concretize`
    lw=None,
    uw=None,
    # common
    clip_neuron_selection_value=-1.0,
    clip_neuron_selection_type="ratio"
):
    """
    If neuron_selection_value >= 0, run an unconstrained/bounds-saving pass
    then a top-K constrained pass; otherwise just one pass.
    """
    # decide which underlying call to use
    def _call_concretize(use_constraints, save_bounds=False, heuristic_indices=None):
        if concretize_mode == 'backward':
            # backward concretize signature
            return backward_concretize(
                self, batch_size, output_dim, lower, upper,
                average_A=average_A,
                node_start=node,
                use_constraints=use_constraints,
                save_bounds=save_bounds,
                heuristic_indices=heuristic_indices,
            )
        elif concretize_mode == 'forward':
            # forward_concretize signature
            return forward_concretize(
                self, lower, upper, lw, uw,
                use_constraints=use_constraints,
                save_bounds=save_bounds,
                heuristic_indices=heuristic_indices,
            )
        else:
            raise ValueError(f"Unknown concretize mode: {concretize_mode}. "
                             "Please use 'backward' or 'forward'.")

    use_constraints = True
    save_bounds = False

    # If clip_neuron_selection_value >= 0, heuristic score-based topk selection is enabled.
    # And we will only apply constrained concretization on topk neurons based on their heuristics.
    # In this case, we'll need to 1) concretize all neurons without any constraints to get a looser bound 
    #                                           --> This is for computing the heuristics
    #                             2) concretize topk neurons with constraints.            
    #                                           --> This is for getting tighter bounds for topk neurons.                
    # In conclusion, if neuron_selection_value >= 0, use_consrtaints will be disabled first.
    # But for the output node in the computational graph we will directly concretize all neurons..
    if clip_neuron_selection_value >= 0 and node.name not in self.output_name:
        use_constraints = False
        # `output_activations` is the list of output activations from current pre-activation node.
        # This output_activations is manually assigned outside of auto_lirpa. Please check 
        #       complete_verifier/input_split/batch_branch_and_bound.py for more info.
        # If a node: 
        #       a) does not have any output_activation, and 
        #       b) heuristic topk selection is enabled, and
        #       c) is not the output node in the computational graph
        #  we will only compute naive bounds on it.
        # Otherwise, we'll need to do both step 1) and 2). And to accelarate step 2), we will save the bounds in 1).
        
        # If 1) this node has at least one output activation node
        #    2) at least one neuron will be selected
        # We will need to concretize with constraints, 
        if node.output_activations is not None and clip_neuron_selection_value > 0:
            save_bounds = True

    # If heuristic topk selection is enabled, this would be the step 1).
    new_lower, new_upper, has_constraints = _call_concretize(
        use_constraints=use_constraints,
        save_bounds=save_bounds,
    )

    # If heuristic topk selection is enabled, this if-branch would be the step 2).
    if (has_constraints
        and node.output_activations is not None
        and clip_neuron_selection_value > 0
        and node.name not in self.output_name):

        score = 0.0
        unstable_masks = False

        # loop through all the output activations to get a comprehensive unstable mask and heuristic score.
        # This output_activations is manually assigned outside of auto_lirpa.
        # Please check complete_verifier/input_split/batch_branch_and_bound.py
        for o_act_node in node.output_activations:
            score = score + o_act_node.compute_bound_improvement_heuristics(new_lower, new_upper)
            unstable_masks = unstable_masks | o_act_node.get_unstable_mask(new_lower, new_upper)
        score = score.flatten(1)                        # shape: (Batchsize, Hidden_dim)
        unstable_masks = unstable_masks.flatten(1)      # shape: (Batchsize, Hidden_dim)

        # Only do second concretize if there exists unstable neurons.
        if unstable_masks.any():
            max_unstable_size = unstable_masks.sum(dim=1).max()
            heuristic_indices = None
            # The K value in topk should be at least 1.
            if clip_neuron_selection_type == "ratio":
                K = max(int(max_unstable_size * clip_neuron_selection_value + 0.5), 1)
            else:
                K = min(clip_neuron_selection_value, max_unstable_size)
            _, heuristic_indices = torch.topk(score, k=K, dim=1, largest=True, sorted=False)
            new_lower, new_upper, _ = _call_concretize(
                use_constraints=True,
                heuristic_indices=heuristic_indices
            )
        else:
            # Previously we've stored to aux bounds, now it should be cleared to avoid any confusion.
            for root in self.roots():
                if (hasattr(root, 'perturbation')
                    and root.perturbation is not None
                    and isinstance(root.perturbation, PerturbationLpNorm)):
                    root.perturbation.clear_aux_bounds()

    return new_lower, new_upper


def concretize_root(self, root, batch_size, output_dim,
                    average_A=False, node_start=None, input_shape=None,
                    use_constraints=False, heuristic_indices=None, save_bounds=False): 
    # The last three optional argument are designed for heuristic-driven constrained concretization.
    # use_constraints:      A flag controling whether to enable constraints solving or not.
    # heuristic_indices:    A index tensor, it select EUQAL number of hidden neurons from each batch. 
    #                           Constrained solving will be further applied on these neurons. Shape (batchsize, n_h_neurons)
    # save_bounds:          A flag determining whether to save naive bounds (to avoid redundant computation)

    if average_A and isinstance(root, BoundParams):
        lA = root.lA.mean(
            node_start.batch_dim + 1, keepdim=True
        ).expand(root.lA.shape) if (root.lA is not None) else None
        uA = root.uA.mean(
            node_start.batch_dim + 1, keepdim=True
        ).expand(root.uA.shape) if (root.uA is not None) else None
    else:
        lA, uA = root.lA, root.uA
    if not isinstance(root.lA, eyeC) and not isinstance(root.lA, Patches):
        lA = root.lA.reshape(output_dim, batch_size, -1).transpose(0, 1) if (lA is not None) else None
    if not isinstance(root.uA, eyeC) and not isinstance(root.uA, Patches):
        uA = root.uA.reshape(output_dim, batch_size, -1).transpose(0, 1) if (uA is not None) else None
    
    has_constraints = False
    if hasattr(root, 'perturbation') and root.perturbation is not None:

        if isinstance(root.perturbation, PerturbationLpNorm):
            # Enable / Disable constraints solving according to `use_constraints`
            root.perturbation.constraints_enable = use_constraints
            if root.perturbation.constraints is not None:
                if self.infeasible_bounds_constraints is not None:
                    root.perturbation.add_infeasible_batches(self.infeasible_bounds_constraints)
                root.perturbation.add_objective_indices(heuristic_indices)
                has_constraints = True

        if isinstance(root, BoundParams):
            # add batch_size dim for weights node
            lb = root.perturbation.concretize(
                root.center.unsqueeze(0), lA, sign=-1, aux=root.aux
            ) if (lA is not None) else None
            ub = root.perturbation.concretize(
                root.center.unsqueeze(0), uA, sign=+1, aux=root.aux
            ) if (uA is not None) else None

        else:
            lb = root.perturbation.concretize(
                root.center, lA, sign=-1, aux=root.aux
            ) if lA is not None else None
            ub = root.perturbation.concretize(
                root.center, uA, sign=+1, aux=root.aux
            ) if uA is not None else None

        if (isinstance(root.perturbation, PerturbationLpNorm) 
            and root.perturbation.constraints is not None
            and root.perturbation.sorted_out_batches["infeasible_batches"] is not None):
            if self.infeasible_bounds_constraints is not None:
                self.infeasible_bounds_constraints = self.infeasible_bounds_constraints | root.perturbation.sorted_out_batches["infeasible_batches"]
            # else:
            #     self.infeasible_bounds_constraints = root.perturbation.sorted_out_batches["infeasible_batches"]

        # If required, save current (naive) bounds to prevent redundant computation next time concretize on the same node
        if isinstance(root.perturbation, PerturbationLpNorm) and root.perturbation.constraints is not None and save_bounds:
            root.perturbation.add_aux_bounds(lb, ub)
        elif isinstance(root.perturbation, PerturbationLpNorm):
        # Otherwise, always clear_aux_bounds to prevent confusion
            root.perturbation.clear_aux_bounds()

    else:
        fv = root.forward_value
        if type(root) == BoundInput:
            # Input node with a batch dimension
            batch_size_ = batch_size
        else:
            # Parameter node without a batch dimension
            batch_size_ = 1

        def concretize_constant(A):
            if isinstance(A, eyeC):
                return fv.view(batch_size_, -1)
            elif isinstance(A, Patches):
                return A.matmul(fv, input_shape=input_shape)
            elif type(root) == BoundInput:
                return A.matmul(fv.view(batch_size_, -1, 1)).squeeze(-1)
            else:
                return A.matmul(fv.view(-1, 1)).squeeze(-1)

        lb = concretize_constant(lA) if (lA is not None) else None
        ub = concretize_constant(uA) if (uA is not None) else None

    return lb, ub, has_constraints


def backward_concretize(self, batch_size, output_dim, lb=None, ub=None,
               average_A=False, node_start=None, 
               use_constraints=False, heuristic_indices=None, save_bounds=False):
    # The last three optional argument are designed for heuristic-driven constrained concretization.
    # use_constraints:      A flag controling whether to enable constraints solving or not.
    # heuristic_indices:    A index tensor, it select EUQAL number of hidden neurons from each batch. 
    #                           Constrained solving will be further applied on these neurons. Shape (batchsize, n_h_neurons)
    # save_bounds:          A flag determining whether to save naive bounds (to avoid redundant computation)
    roots = self.roots()
    if isinstance(lb, torch.Tensor) and lb.ndim > 2:
        lb = lb.reshape(lb.shape[0], -1)
    if isinstance(ub, torch.Tensor) and ub.ndim > 2:
        ub = ub.reshape(ub.shape[0], -1)

    def add_b(b1, b2):
        if b2 is None:
            return b1
        elif b1 is None:
            return b2
        # Check if b1 is a tensor and if all its elements are infinity
        if torch.is_tensor(b1) and torch.isinf(b1).all():
            return b1
        # Check if b2 is a tensor and if all its elements are infinity
        if torch.is_tensor(b2) and torch.isinf(b2).all():
            return b2
        else:
            return b1 + b2

    has_constraints = False
    for root in roots:
        root.lb = root.ub = None
        if root.lA is None and root.uA is None:
            continue
        root.lb, root.ub, has_constraints_this_root = self.concretize_root(
            root, batch_size, output_dim, average_A=average_A,
            node_start=node_start, input_shape=roots[0].center.shape,
            use_constraints=use_constraints, heuristic_indices=heuristic_indices, save_bounds=save_bounds)

        has_constraints = has_constraints | has_constraints_this_root

        lb = add_b(lb, root.lb)
        ub = add_b(ub, root.ub)

    return lb, ub, has_constraints


def forward_concretize(self, lower, upper, lw, uw, use_constraints=False, heuristic_indices=None, save_bounds=False):
    """
    Concretize function for forward bound. 

    :param lower:                   Tensor. Intermediate layer lower bounds.
    :param upper:                   Tensor. Intermediate layer upper bounds.
    :param lw:                      Tensor. Intermediate layer lower A matrix.
    :param uw:                      Tensor. Intermediate layer upper A matrix.
    :param use_constraints:         bool. A flag controling whether to enable constraints solving or not.
        If heuristic ratio is set, the first concretization run should disbale constraints solving.
    :param heuristic_indices:       Index Tensor. A index tensor, it select **equal** number of hidden neurons from each batch.
        Constrained solving will be further applied on these neurons. Shape (batchsize, n_h_neurons)
    :param save_bounds:             bool. A flag controling whether to save naive bounds.
    
    :return res_lower:              Tensor. The lower bound tensor.
    :return res_upper:              Tensor. The upper bound tensor.
    :return has_constraints:        bool. Whether constraints has been stored.
    """
    res_lower = 0.0
    res_upper = 0.0
    prev_dim_in = 0
    has_constraints = False
    roots = self.roots()
    assert (lw.ndim > 1)
    lA = lw.reshape(self.batch_size, self.dim_in, -1).transpose(1, 2)
    uA = uw.reshape(self.batch_size, self.dim_in, -1).transpose(1, 2)
    for root in roots:
        if hasattr(root, 'perturbation') and root.perturbation is not None:
            _lA = lA[:, :, prev_dim_in : (prev_dim_in + root.dim)]
            _uA = uA[:, :, prev_dim_in : (prev_dim_in + root.dim)]

            if isinstance(root.perturbation, PerturbationLpNorm):
                root.perturbation.constraints_enable = use_constraints
                if root.perturbation.constraints is not None:
                    if self.infeasible_bounds_constraints is not None:
                        root.perturbation.add_infeasible_batches(self.infeasible_bounds_constraints)
                    root.perturbation.add_objective_indices(heuristic_indices)
                    has_constraints = True                 

            # Previously added concretized bounds directly to lower/upper.
            # Now extract them first for reuse (e.g., in aux_bounds).
            temp_lower = root.perturbation.concretize(
                root.center, _lA, sign=-1, aux=root.aux
                ).view(lower.shape)
            temp_upper = root.perturbation.concretize(
                root.center, _uA, sign=+1, aux=root.aux
                ).view(upper.shape)
            
            # Update infeasible_batches
            if (isinstance(root.perturbation, PerturbationLpNorm)
                and root.perturbation.constraints is not None 
                and root.perturbation.sorted_out_batches["infeasible_batches"] is not None):
                if self.infeasible_bounds_constraints is not None:
                    self.infeasible_bounds_constraints = self.infeasible_bounds_constraints | root.perturbation.sorted_out_batches["infeasible_batches"]
                # else:
                #     self.infeasible_bounds_constraints = root.perturbation.sorted_out_batches["infeasible_batches"]

            # If required, save current (naive) bounds to prevent redundant computation next time concretize on the same node
            if isinstance(root.perturbation, PerturbationLpNorm) and root.perturbation.constraints is not None and save_bounds:
                root.perturbation.add_aux_bounds(temp_lower, temp_upper)
            elif isinstance(root.perturbation, PerturbationLpNorm):
            # Otherwise, always clear_aux_bounds to prevent confusion
                root.perturbation.clear_aux_bounds()

            # Now the concretization result from this root will be accumulated into final bounds.
            # Here we add temp_lower onto res_lower, instead of lower. 
            # It's because the lower value will be used twice, any modification to it should be avoided.
            res_lower = res_lower + temp_lower
            res_upper = res_upper + temp_upper                        
    
    res_lower = res_lower + lower
    res_upper = res_upper + upper
    return res_lower, res_upper, has_constraints


================================================
FILE: auto_LiRPA/concretize_func.py
================================================
#########################################################################
##   This file is part of the auto_LiRPA library, a core part of the   ##
##   α,β-CROWN (alpha-beta-CROWN) neural network verifier developed    ##
##   by the α,β-CROWN Team                                             ##
##                                                                     ##
##   Copyright (C) 2020-2025 The α,β-CROWN Team                        ##
##   Team leaders:                                                     ##
##          Faculty:   Huan Zhang <huan@huan-zhang.com> (UIUC)         ##
##          Student:   Xiangru Zhong <xiangru4@illinois.edu> (UIUC)    ##
##                                                                     ##
##   See CONTRIBUTORS for all current and past developers in the team. ##
##                                                                     ##
##     This program is licensed under the BSD 3-Clause License,        ##
##        contained in the LICENCE file in this directory.             ##
##                                                                     ##
#########################################################################
import torch

from math import floor, ceil
from .utils import eyeC

# Declaration of the shape naming:

# B / batchsize  : The number of batches. In this `concretize_func.py`, if a tensor has batch dimension, we assume
#                    it will only be the first dimention of this tensor . That is: B = tensor.shape[0]
# 
# B_act          : The number of active batches. We will only apply constraints to a subset of batches, and these
#                   batches are called active batches. B_act <= B. There are two cases:
#                       -- When `no_return_inf` mode is disabled, we will keep B_act static throughout the entire 
#                           BaB iteration. In this case, B_act equals the number of batches not fully covered by 
#                           constraints, as determined by `sort_out_constr_batches` function.
#                       -- When `no_return_inf` mode is enabled, then B_act decreases over iterations, since more
#                           batches will be marked as infeasible. See `PerturbationLpNorm.add_infeasible_batches`.
# 
# X / x_dim      : The number of input neurons (batch dimension excluded). It stands for the input shape of the
#                   neural network. For tensors such as x0, epsilon, x_U, x_L, X = prod(*tensor.shape[1:])
# 
# H / hidden_dim : The number of hidden neurons (batch and input dimension excluded). It stands for the output
#                   shape of this hidden layer. For the objective A tensor, there are two cases:
#                       -- The tensor has batch dimention: H = tensor.view(B, -1, X).shape[1]
#                       -- The tensor does not have batch dimention: H = tensor.view(-1, X).shape[0]
# 
# H_act          : The number of active batches. We may only apply constraints to a subset of hidden neurons,
#                   and these neurons are called active neurons. H_act <= H.
#
# N_constr       : The number of constraints. For constraints_A matrix:
#                       -- In `sort_out_constr_batches` function, its shape is (B, N_constr, X)
#                       -- In `constraints_solving` function, its shape is (B_act, N_constt, X)

def construct_constraints(constr_A: torch.Tensor, constr_b: torch.Tensor, constr_rhs: torch.Tensor,
                            batchsize, x_dim, sign=1):
    r"""
    Construct the constraints tuple. This function provides a unified interface to generate this tuple.
    All the users should carefully read this function to fully understand the standard form of constraints.

    The first three argument expresses the non-standard form of the constraints:
                                    A @ x + b <= rhs
    We will first convert it into the standard form:
                                    A @ x + b' <= 0
    The the standard expression of constraints should be (constr_A, constr_b')

    Args:
        constr_A:   The coefficient A matrix of constraints.
                        It should be able to be reshaped into: (B, N_constr, X)
        constr_b:   The bias term of constraints.        
                        It should be able to be reshaped into: (B, N_constr)
        constr_rhs: The right-hand-side term of constraints.        
                        It should be able
Download .txt
gitextract_yr8n0etx/

├── .github/
│   └── ISSUE_TEMPLATE/
│       └── bug_report.md
├── .gitignore
├── .readthedocs.yaml
├── CONTRIBUTORS
├── LICENSE
├── README.md
├── auto_LiRPA/
│   ├── __init__.py
│   ├── backward_bound.py
│   ├── beta_crown.py
│   ├── bound_general.py
│   ├── bound_multi_gpu.py
│   ├── bound_op_map.py
│   ├── bound_ops.py
│   ├── bounded_tensor.py
│   ├── concretize_bounds.py
│   ├── concretize_func.py
│   ├── cuda/
│   │   ├── cuda_kernels.cu
│   │   └── cuda_utils.cpp
│   ├── cuda_utils.py
│   ├── edit_graph.py
│   ├── eps_scheduler.py
│   ├── forward_bound.py
│   ├── interval_bound.py
│   ├── jacobian.py
│   ├── linear_bound.py
│   ├── operators/
│   │   ├── __init__.py
│   │   ├── activation_base.py
│   │   ├── activations.py
│   │   ├── add_sub.py
│   │   ├── base.py
│   │   ├── bivariate.py
│   │   ├── clampmult.py
│   │   ├── constant.py
│   │   ├── convex_concave.py
│   │   ├── convolution.py
│   │   ├── cut_ops.py
│   │   ├── dropout.py
│   │   ├── dtype.py
│   │   ├── gelu.py
│   │   ├── indexing.py
│   │   ├── jacobian.py
│   │   ├── leaf.py
│   │   ├── linear.py
│   │   ├── logical.py
│   │   ├── minmax.py
│   │   ├── normalization.py
│   │   ├── pooling.py
│   │   ├── reduce.py
│   │   ├── relu.py
│   │   ├── reshape.py
│   │   ├── resize.py
│   │   ├── rnn.py
│   │   ├── s_shaped.py
│   │   ├── shape.py
│   │   ├── slice_concat.py
│   │   ├── softmax.py
│   │   ├── solver_utils.py
│   │   ├── tile.py
│   │   └── trigonometric.py
│   ├── opt_pruner.py
│   ├── optimize_graph.py
│   ├── optimized_bounds.py
│   ├── output_constraints.py
│   ├── parse_graph.py
│   ├── patches.py
│   ├── perturbations.py
│   ├── solver_module.py
│   ├── tools.py
│   ├── utils.py
│   └── wrapper.py
├── doc/
│   ├── .gitignore
│   ├── Makefile
│   ├── README.md
│   ├── api.rst
│   ├── conf.py
│   ├── index.rst
│   └── process.py
├── examples/
│   ├── .gitignore
│   ├── __init__.py
│   ├── language/
│   │   ├── .gitignore
│   │   ├── Transformer/
│   │   │   ├── Transformer.py
│   │   │   ├── __init__.py
│   │   │   ├── modeling.py
│   │   │   └── utils.py
│   │   ├── data_utils.py
│   │   ├── language_utils.py
│   │   ├── lstm.py
│   │   ├── oracle.py
│   │   ├── preprocess/
│   │   │   ├── pre_compute_lm_scores.py
│   │   │   └── preprocess_sst.py
│   │   └── train.py
│   ├── sequence/
│   │   ├── .gitignore
│   │   ├── __init__.py
│   │   ├── data_utils.py
│   │   ├── lstm.py
│   │   └── train.py
│   ├── simple/
│   │   ├── invprop.py
│   │   ├── lp_full.py
│   │   ├── mip_lp_solver.py
│   │   ├── models/
│   │   │   └── spectral_NOR_MLP_B.pth
│   │   └── toy.py
│   └── vision/
│       ├── .gitignore
│       ├── bound_option.py
│       ├── cifar_training.py
│       ├── custom_op.py
│       ├── data/
│       │   ├── .gitignore
│       │   ├── ImageNet64/
│       │   │   └── imagenet_data_loader.py
│       │   └── tinyImageNet/
│       │       ├── .gitignore
│       │       └── tinyimagenet_download.sh
│       ├── datasets.py
│       ├── efficient_convolution.py
│       ├── imagenet_training.py
│       ├── jacobian.py
│       ├── models/
│       │   ├── __init__.py
│       │   ├── densenet.py
│       │   ├── densenet_imagenet.py
│       │   ├── densenet_no_bn.py
│       │   ├── feedforward.py
│       │   ├── mobilenet.py
│       │   ├── resnet.py
│       │   ├── resnet18.py
│       │   ├── resnext.py
│       │   ├── resnext_imagenet64.py
│       │   ├── vnncomp_resnet.py
│       │   ├── wide_resnet_cifar.py
│       │   └── wide_resnet_imagenet64.py
│       ├── pretrained/
│       │   ├── cifar_2c2f.pth
│       │   ├── kw_mnist.pth
│       │   ├── mnist_a_adv.pth
│       │   ├── mnist_cnn_small.pth
│       │   ├── mnist_fc_3layer.pth
│       │   └── test_min_max.pth
│       ├── save_intermediate_bound.py
│       ├── simple_training.py
│       ├── simple_verification.py
│       ├── tinyimagenet_training.py
│       ├── verify_two_node.py
│       └── weight_perturbation_training.py
├── setup.py
└── tests/
    ├── .gitignore
    ├── data/
    │   ├── .gitignore
    │   ├── avgpool_test_data
    │   ├── beta_crown_test_data
    │   ├── bound_ops_data
    │   ├── ckpt_lstm
    │   ├── ckpt_transformer
    │   ├── constant_test_data
    │   ├── conv1d_test_data_3-0-2
    │   ├── conv1d_test_data_3-0-3
    │   ├── conv1d_test_data_3-1-2
    │   ├── conv1d_test_data_3-1-3
    │   ├── conv1d_test_data_4-0-2
    │   ├── conv1d_test_data_4-0-3
    │   ├── conv1d_test_data_4-1-2
    │   ├── conv1d_test_data_4-1-3
    │   ├── distinct_patches_test_data
    │   ├── invprop/
    │   │   ├── ood.onnx
    │   │   ├── ood_reference
    │   │   └── simple_reference
    │   ├── jacobian_test_data
    │   ├── language_test_data
    │   ├── maxpool_test_data_3-0-3-0
    │   ├── maxpool_test_data_3-0-3-1
    │   ├── maxpool_test_data_3-1-3-0
    │   ├── maxpool_test_data_3-1-3-1
    │   ├── maxpool_test_data_4-0-4-0
    │   ├── maxpool_test_data_4-0-4-1
    │   ├── maxpool_test_data_4-1-4-0
    │   ├── maxpool_test_data_4-1-4-1
    │   ├── min_max_test_data
    │   ├── rectangle_patches_test_data
    │   ├── resnet_patches_test_data
    │   ├── s_shape_test_data
    │   ├── test_constrained_concretize
    │   ├── test_general_shape_data
    │   ├── test_perturbation_data
    │   ├── test_save_data
    │   ├── vision_clip_test_data
    │   ├── vision_test_data
    │   └── weight_perturbation_test_data
    ├── data_64/
    │   ├── avgpool_test_data
    │   ├── bound_ops_data
    │   ├── constant_test_data
    │   ├── conv1d_test_data_3-0-2
    │   ├── conv1d_test_data_3-0-3
    │   ├── conv1d_test_data_3-1-2
    │   ├── conv1d_test_data_3-1-3
    │   ├── conv1d_test_data_4-0-2
    │   ├── conv1d_test_data_4-0-3
    │   ├── conv1d_test_data_4-1-2
    │   ├── conv1d_test_data_4-1-3
    │   ├── general_shape_data
    │   ├── invprop/
    │   │   ├── ood_reference
    │   │   └── simple_reference
    │   ├── jacobian_test_data
    │   ├── maxpool_test_data_3-0-3-0
    │   ├── maxpool_test_data_3-0-3-1
    │   ├── maxpool_test_data_3-1-3-0
    │   ├── maxpool_test_data_3-1-3-1
    │   ├── maxpool_test_data_4-0-4-0
    │   ├── maxpool_test_data_4-0-4-1
    │   ├── maxpool_test_data_4-1-4-0
    │   ├── maxpool_test_data_4-1-4-1
    │   ├── min_max_test_data
    │   ├── rectangle_patches_test_data
    │   ├── resnet_patches_test_data
    │   ├── s_shape_test_data
    │   ├── test_constrained_concretize
    │   ├── test_general_shape_data
    │   ├── test_save_data
    │   ├── vision_clip_test_data
    │   ├── vision_test_data
    │   └── weight_perturbation_test_data
    ├── test_1d_activation.py
    ├── test_2d_activation.py
    ├── test_avgpool.py
    ├── test_bound_ops.py
    ├── test_branching_heuristics.py
    ├── test_clip_domains.py
    ├── test_constant.py
    ├── test_constrained_concretize.py
    ├── test_conv.py
    ├── test_conv1d.py
    ├── test_distinct_patches.py
    ├── test_examples.py
    ├── test_examples_ci.py
    ├── test_general_nonlinear.py
    ├── test_general_shape.py
    ├── test_identity.py
    ├── test_invprop.py
    ├── test_jacobian.py
    ├── test_language_models.py
    ├── test_linear_cnn_model.py
    ├── test_linear_model.py
    ├── test_maxpool.py
    ├── test_min_max.py
    ├── test_perturbation.py
    ├── test_rectangle_patches.py
    ├── test_resnet_patches.py
    ├── test_s_shaped.py
    ├── test_save_intermediate.py
    ├── test_simple_verification.py
    ├── test_state_dict_name.py
    ├── test_tensor_storage.py
    ├── test_upsample.py
    ├── test_vision_models.py
    ├── test_vision_models_hardtanh.py
    ├── test_weight_perturbation.py
    └── testcase.py
Download .txt
SYMBOL INDEX (1697 symbols across 139 files)

FILE: auto_LiRPA/backward_bound.py
  function batched_backward (line 32) | def batched_backward(self: 'BoundedModule', node, C, unstable_idx, batch...
  function backward_general (line 127) | def backward_general(
  function get_unstable_size (line 398) | def get_unstable_size(unstable_idx):
  function check_optimized_variable_sparsity (line 405) | def check_optimized_variable_sparsity(self: 'BoundedModule', node):
  function get_sparse_C (line 423) | def get_sparse_C(self: 'BoundedModule', node, ref_intermediate):
  function restore_sparse_bounds (line 608) | def restore_sparse_bounds(self: 'BoundedModule', node, unstable_idx,
  function get_degrees (line 643) | def get_degrees(node_start):
  function _preprocess_C (line 662) | def _preprocess_C(self: 'BoundedModule', C, node):
  function addA (line 708) | def addA(A1, A2):
  function add_bound (line 720) | def add_bound(node, node_pre, lA=None, uA=None):
  function add_constant_node (line 749) | def add_constant_node(lb, ub, node):
  function save_root_A (line 761) | def save_root_A(node, A_record, A_dict, roots, needed_A_dict, lb, ub,
  function select_unstable_idx (line 801) | def select_unstable_idx(ref_intermediate_lb, ref_intermediate_ub, unstab...
  function get_unstable_locations (line 814) | def get_unstable_locations(self: 'BoundedModule', ref_intermediate_lb,
  function get_alpha_crown_start_nodes (line 856) | def get_alpha_crown_start_nodes(
  function merge_A (line 949) | def merge_A(node, batch_A, ret_A):

FILE: auto_LiRPA/beta_crown.py
  class SparseBeta (line 28) | class SparseBeta:
    method __init__ (line 29) | def __init__(self, shape, bias=False, betas=None, device='cpu'):
    method apply_splits (line 42) | def apply_splits(self, history, key):
  function get_split_nodes (line 62) | def get_split_nodes(self: 'BoundedModule'):
  function set_beta (line 82) | def set_beta(self: 'BoundedModule', enable_opt_interm_bounds, parameters,
  function set_beta_cuts (line 121) | def set_beta_cuts(self: 'BoundedModule', parameters, lr_cut_beta, betas,
  function reset_beta (line 132) | def reset_beta(self: 'BoundedModule', node, shape, betas, bias=False,
  function beta_crown_backward_bound (line 150) | def beta_crown_backward_bound(self: 'BoundedModule', node, lA, uA, start...
  function print_optimized_beta (line 246) | def print_optimized_beta(acts):

FILE: auto_LiRPA/bound_general.py
  class BoundedModule (line 40) | class BoundedModule(nn.Module):
    method __init__ (line 62) | def __init__(self, model, global_input, bound_opts=None,
    method nodes (line 192) | def nodes(self) -> List[Bound]:
    method get_enabled_opt_act (line 195) | def get_enabled_opt_act(self):
    method get_optimizable_activations (line 202) | def get_optimizable_activations(self):
    method get_perturbed_optimizable_activations (line 218) | def get_perturbed_optimizable_activations(self):
    method get_splittable_activations (line 221) | def get_splittable_activations(self):
    method get_layers_requiring_bounds (line 225) | def get_layers_requiring_bounds(self):
    method check_incompatible_nodes (line 254) | def check_incompatible_nodes(self, model):
    method non_deter_wrapper (line 283) | def non_deter_wrapper(self, op, *args, **kwargs):
    method non_deter_scatter_add (line 293) | def non_deter_scatter_add(self, *args, **kwargs):
    method non_deter_index_select (line 296) | def non_deter_index_select(self, *args, **kwargs):
    method set_bound_opts (line 299) | def set_bound_opts(self, new_opts):
    method set_gcp_relu_indicators (line 307) | def set_gcp_relu_indicators(self, relu_layer_name, relu_indicators):
    method _get_A_norm (line 327) | def _get_A_norm(A):
    method __call__ (line 340) | def __call__(self, *input, **kwargs):
    method register_parameter (line 356) | def register_parameter(self, name, param):
    method _named_members (line 393) | def _named_members(self,
    method train (line 416) | def train(self, mode=True):
    method eval (line 421) | def eval(self):
    method to (line 426) | def to(self, *args, **kwargs):
    method __getitem__ (line 444) | def __getitem__(self, name):
    method roots (line 450) | def roots(self):
    method final_node (line 453) | def final_node(self):
    method get_forward_value (line 456) | def get_forward_value(self, node):
    method forward (line 479) | def forward(self, *x, final_node_name=None,
    method _mark_perturbed_nodes (line 515) | def _mark_perturbed_nodes(self, input):
    method _check_patches_mode (line 570) | def _check_patches_mode(self):
    method _clear_and_set_new (line 587) | def _clear_and_set_new(
    method set_input (line 637) | def set_input(
    method _get_node_input (line 665) | def _get_node_input(self, nodesOP, nodesIn, node):
    method _to (line 682) | def _to(self, obj, dest, inplace=False):
    method _convert_nodes (line 705) | def _convert_nodes(self, model, global_input):
    method _build_graph (line 800) | def _build_graph(self, nodesOP, nodesIn, nodesOut, template):
    method rename_nodes (line 814) | def rename_nodes(self, nodesOP, nodesIn, rename_dict):
    method _split_complex (line 823) | def _split_complex(self, nodesOP, nodesIn):
    method _get_node_name_map (line 872) | def _get_node_name_map(self):
    method _convert (line 889) | def _convert(self, model, global_input):
    method check_prior_bounds (line 923) | def check_prior_bounds(self, node, C=None):
    method compute_intermediate_bounds (line 970) | def compute_intermediate_bounds(self, node: Bound, prior_checked=False):
    method get_ref_intermediate_bounds (line 1099) | def get_ref_intermediate_bounds(self, node):
    method merge_A_dict (line 1131) | def merge_A_dict(self, lA_dict, uA_dict):
    method compute_bounds (line 1146) | def compute_bounds(
    method save_intermediate (line 1427) | def save_intermediate(self, save_path=None):
    method _compute_bounds_main (line 1461) | def _compute_bounds_main(self, C=None, method='backward', IBP=False,
    method _set_used_nodes (line 1543) | def _set_used_nodes(self, final):
    method init_infeasible_bounds_constraints (line 1560) | def init_infeasible_bounds_constraints(self, batchsize, device):

FILE: auto_LiRPA/bound_multi_gpu.py
  class BoundDataParallel (line 22) | class BoundDataParallel(DataParallel):
    method __init__ (line 25) | def __init__(self, *inputs, **kwargs):
    method forward (line 30) | def forward(self, *inputs, **kwargs):
    method get_property (line 122) | def get_property(model, node_class=None, att_name=None, node_name=None):
    method state_dict (line 137) | def state_dict(self, destination=None, prefix='', keep_vars=False):
    method _named_members (line 141) | def _named_members(self, get_members_fn, prefix='', recurse=True, remo...
    method __getitem__ (line 144) | def __getitem__(self, name):

FILE: auto_LiRPA/bound_op_map.py
  function register_custom_op (line 34) | def register_custom_op(op_name: str, bound_obj: Bound) -> None:
  function unregister_custom_op (line 37) | def unregister_custom_op(op_name: str) -> None:

FILE: auto_LiRPA/bounded_tensor.py
  class BoundedTensor (line 23) | class BoundedTensor(Tensor):
    method __new__ (line 26) | def __new__(cls, x, ptb=None, *args, **kwargs):
    method __init__ (line 35) | def __init__(self, x, ptb=None):
    method __repr__ (line 38) | def __repr__(self):
    method clone (line 44) | def clone(self, *args, **kwargs):
    method _func (line 48) | def _func(self, func, *args, **kwargs):
    method to (line 56) | def to(self, *args, **kwargs):
    method _convert (line 67) | def _convert(cls, ret):
    method __torch_function__ (line 85) | def __torch_function__(cls, func, types, args=(), kwargs=None):
  class BoundedParameter (line 97) | class BoundedParameter(nn.Parameter):
    method __new__ (line 98) | def __new__(cls, data, ptb, requires_grad=True):
    method __init__ (line 101) | def __init__(self, data, ptb, requires_grad=True):
    method __deepcopy__ (line 105) | def __deepcopy__(self, memo):
    method __repr__ (line 113) | def __repr__(self):
    method __reduce_ex__ (line 117) | def __reduce_ex__(self, proto):

FILE: auto_LiRPA/concretize_bounds.py
  function concretize_bounds (line 29) | def concretize_bounds(
  function concretize_root (line 152) | def concretize_root(self, root, batch_size, output_dim,
  function backward_concretize (line 244) | def backward_concretize(self, batch_size, output_dim, lb=None, ub=None,
  function forward_concretize (line 290) | def forward_concretize(self, lower, upper, lw, uw, use_constraints=False...

FILE: auto_LiRPA/concretize_func.py
  function construct_constraints (line 50) | def construct_constraints(constr_A: torch.Tensor, constr_b: torch.Tensor...
  function _sort_out_constraints (line 79) | def _sort_out_constraints(A, b, x0, epsilon):
  function _dist_rearrange (line 110) | def _dist_rearrange(constraints_A, constraints_b, x0):
  function _solve_dual_var (line 136) | def _solve_dual_var(constr_a, object_a, constr_d, epsilon, a_mul_e=None):
  function sort_out_constr_batches (line 220) | def sort_out_constr_batches(x_L, x_U, constraints, rearrange_constraints...
  function constraints_solving (line 303) | def constraints_solving(

FILE: auto_LiRPA/cuda/cuda_utils.cpp
  function double2float_foward (line 10) | torch::Tensor double2float_foward(
  function PYBIND11_MODULE (line 23) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {

FILE: auto_LiRPA/cuda_utils.py
  class DummyCudaClass (line 23) | class DummyCudaClass:
    method __getattr__ (line 25) | def __getattr__(self, attr):
  function test_double2float (line 76) | def test_double2float():

FILE: auto_LiRPA/edit_graph.py
  function add_nodes (line 27) | def add_nodes(self: 'BoundedModule', nodes):
  function add_input_node (line 46) | def add_input_node(self: 'BoundedModule', node, index=None):
  function delete_node (line 56) | def delete_node(self: 'BoundedModule', node):
  function replace_node (line 68) | def replace_node(self: 'BoundedModule', node_old, node_new):

FILE: auto_LiRPA/eps_scheduler.py
  class BaseScheduler (line 20) | class BaseScheduler(object):
    method __init__ (line 21) | def __init__(self, max_eps, opt_str):
    method __repr__ (line 30) | def __repr__(self):
    method parse_opts (line 33) | def parse_opts(self, s):
    method get_max_eps (line 41) | def get_max_eps(self):
    method get_eps (line 44) | def get_eps(self):
    method reached_max_eps (line 47) | def reached_max_eps(self):
    method step_batch (line 50) | def step_batch(self, verbose=False):
    method step_epoch (line 55) | def step_epoch(self, verbose=False):
    method update_loss (line 60) | def update_loss(self, new_loss):
    method train (line 64) | def train(self):
    method eval (line 67) | def eval(self):
    method set_epoch_length (line 71) | def set_epoch_length(self, epoch_length):
  class FixedScheduler (line 75) | class FixedScheduler(BaseScheduler):
    method __init__ (line 76) | def __init__(self, max_eps, opt_str=""):
  class LinearScheduler (line 81) | class LinearScheduler(BaseScheduler):
    method __init__ (line 83) | def __init__(self, max_eps, opt_str):
    method __repr__ (line 89) | def __repr__(self):
    method step_epoch (line 93) | def step_epoch(self, verbose = True):
    method step_batch (line 111) | def step_batch(self):
  class RangeScheduler (line 120) | class RangeScheduler(BaseScheduler):
    method __init__ (line 122) | def __init__(self, max_eps, opt_str):
    method __repr__ (line 127) | def __repr__(self):
    method step_epoch (line 131) | def step_epoch(self, verbose = True):
    method step_batch (line 138) | def step_batch(self):
  class BiLinearScheduler (line 141) | class BiLinearScheduler(LinearScheduler):
    method __init__ (line 143) | def __init__(self, max_eps, opt_str):
    method __repr__ (line 150) | def __repr__(self):
    method step_epoch (line 154) | def step_epoch(self, verbose = True):
  class SmoothedScheduler (line 175) | class SmoothedScheduler(BaseScheduler):
    method __init__ (line 177) | def __init__(self, max_eps, opt_str):
    method set_epoch_length (line 193) | def set_epoch_length(self, epoch_length):
    method step_epoch (line 200) | def step_epoch(self, verbose = True):
    method step_batch (line 210) | def step_batch(self, verbose=False):
  class AdaptiveScheduler (line 236) | class AdaptiveScheduler(BaseScheduler):
    method __init__ (line 237) | def __init__(self, max_eps, opt_str):
    method step_batch (line 250) | def step_batch(self):

FILE: auto_LiRPA/forward_bound.py
  function forward_general (line 32) | def forward_general(self: 'BoundedModule', C=None, node:'Bound'=None, co...
  function forward_general_dynamic (line 119) | def forward_general_dynamic(self: 'BoundedModule', C=None, node:'Bound'=...
  function clean_memory (line 241) | def clean_memory(self: 'BoundedModule', node):
  function forward_refinement (line 258) | def forward_refinement(self: 'BoundedModule', node):
  function init_forward (line 278) | def init_forward(self: 'BoundedModule', roots, dim_in):

FILE: auto_LiRPA/interval_bound.py
  function IBP_general (line 26) | def IBP_general(self: 'BoundedModule', node=None, C=None,
  function _IBP_loss_fusion (line 104) | def _IBP_loss_fusion(self: 'BoundedModule', node, C):
  function check_IBP_intermediate (line 152) | def check_IBP_intermediate(self: 'BoundedModule', node):
  function check_IBP_first_linear (line 198) | def check_IBP_first_linear(self: 'BoundedModule', node):
  function compare_with_IBP (line 233) | def compare_with_IBP(self, node, lower, upper, C=None):

FILE: auto_LiRPA/jacobian.py
  function _expand_jacobian (line 32) | def _expand_jacobian(self):
  function expand_jacobian_node (line 57) | def expand_jacobian_node(self, jacobian_node):
  function compute_jacobian_bounds (line 177) | def compute_jacobian_bounds(self: 'BoundedModule', x, optimize=True,

FILE: auto_LiRPA/linear_bound.py
  class LinearBound (line 17) | class LinearBound:
    method __init__ (line 18) | def __init__(
    method is_single_bound (line 40) | def is_single_bound(self):

FILE: auto_LiRPA/operators/activation_base.py
  class BoundActivation (line 28) | class BoundActivation(Bound):
    method __init__ (line 29) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method _init_masks (line 38) | def _init_masks(self, x):
    method init_linear_relaxation (line 43) | def init_linear_relaxation(self, x):
    method add_linear_relaxation (line 50) | def add_linear_relaxation(self, mask, type, k, x0, y0=None):
    method bound_relax (line 81) | def bound_relax(self, x, init=False):
    method bound_backward (line 84) | def bound_backward(self, last_lA, last_uA, x, reduce_bias=True, **kwar...
    method bound_forward_w (line 127) | def bound_forward_w(
    method bound_forward_b (line 137) | def bound_forward_b(
    method bound_forward (line 144) | def bound_forward(self, dim_in, x):
    method interval_propagate (line 161) | def interval_propagate(self, *v):
    method get_split_mask (line 165) | def get_split_mask(self, lower, upper, input_index):
    method compute_bound_improvement_heuristics (line 173) | def compute_bound_improvement_heuristics(self, lower, upper):
    method get_unstable_mask (line 180) | def get_unstable_mask(self, lower, upper):
  class BoundOptimizableActivation (line 188) | class BoundOptimizableActivation(BoundActivation):
    method __init__ (line 189) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method opt_init (line 208) | def opt_init(self):
    method opt_start (line 213) | def opt_start(self):
    method opt_reuse (line 217) | def opt_reuse(self):
    method opt_no_reuse (line 221) | def opt_no_reuse(self):
    method opt_end (line 226) | def opt_end(self):
    method clip_alpha (line 230) | def clip_alpha(self):
    method init_opt_parameters (line 233) | def init_opt_parameters(self, start_nodes):
    method _init_opt_parameters_impl (line 244) | def _init_opt_parameters_impl(self, size_spec, name_start=None):
    method init_linear_relaxation (line 248) | def init_linear_relaxation(self, x, dim_opt=None):
    method bound_relax (line 275) | def bound_relax(self, x, init=False, dim_opt=None):
    method bound_backward (line 278) | def bound_backward(self, last_lA, last_uA, x, start_node=None,
    method _no_bound_parameters (line 337) | def _no_bound_parameters(self):
    method _transfer_alpha (line 342) | def _transfer_alpha(self, alpha, device=None, dtype=None, non_blocking...
    method dump_alpha (line 347) | def dump_alpha(self, device=None, dtype=None, non_blocking=False):
    method restore_alpha (line 353) | def restore_alpha(self, alpha, device=None, dtype=None, non_blocking=F...
    method drop_unused_alpha (line 359) | def drop_unused_alpha(self, keep_nodes):

FILE: auto_LiRPA/operators/activations.py
  class BoundSoftplus (line 29) | class BoundSoftplus(BoundActivation):
    method __init__ (line 30) | def __init__(self, *args, **kwargs):
    method forward (line 34) | def forward(self, x):
  class BoundAbs (line 38) | class BoundAbs(BoundActivation):
    method __init__ (line 39) | def __init__(self, *args, **kwargs):
    method forward (line 43) | def forward(self, x):
    method bound_relax (line 46) | def bound_relax(self, x, init=False):
    method bound_backward (line 58) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method interval_propagate (line 98) | def interval_propagate(self, *v):
  class BoundATenHeaviside (line 105) | class BoundATenHeaviside(BoundOptimizableActivation):
    method forward (line 106) | def forward(self, *x):
    method interval_propagate (line 111) | def interval_propagate(self, *v):
    method _init_opt_parameters_impl (line 115) | def _init_opt_parameters_impl(self, size_spec, name_start):
    method clip_alpha (line 120) | def clip_alpha(self):
    method bound_backward (line 124) | def bound_backward(self, last_lA, last_uA, *x, start_node=None,
  class BoundSqr (line 170) | class BoundSqr(BoundOptimizableActivation):
    method __init__ (line 172) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 176) | def forward(self, x):
    method bound_relax (line 179) | def bound_relax(self, x, init=False, dim_opt=None):
    method _init_opt_parameters_impl (line 200) | def _init_opt_parameters_impl(self, size_spec, **kwargs):
    method interval_propagate (line 207) | def interval_propagate(self, *v):
    method build_gradient_node (line 213) | def build_gradient_node(self, grad_upstream):
  class SqrGrad (line 217) | class SqrGrad(Module):
    method forward (line 218) | def forward(self, grad_last, preact):
  class BoundHardTanh (line 223) | class BoundHardTanh(BoundActivation):
    method __init__ (line 225) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 232) | def forward(self, x, min_val, max_val):
    method bound_backward (line 235) | def bound_backward(self, last_lA, last_uA, x, min_val, max_val, start_...
    method bound_relax (line 275) | def bound_relax(self, x, min_val, max_val, init=False, dim_opt=None):
    method interval_propagate (line 405) | def interval_propagate(self, *v):
  class BoundFloor (line 413) | class BoundFloor(BoundActivation):
    method forward (line 414) | def forward(self, x):
    method bound_relax (line 417) | def bound_relax(self, x, init=False):
  class BoundMultiPiecewiseNonlinear (line 424) | class BoundMultiPiecewiseNonlinear(BoundOptimizableActivation):
    method __init__ (line 425) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 429) | def forward(self, x, weight, offset):
    method clip_alpha (line 432) | def clip_alpha(self):
    method bound_backward (line 436) | def bound_backward(self, last_lA, last_uA, x, weight, offset,
    method _init_opt_parameters_impl (line 493) | def _init_opt_parameters_impl(self, size_spec, **kwargs):
    method get_split_mask (line 499) | def get_split_mask(self, lower, upper, input_index):

FILE: auto_LiRPA/operators/add_sub.py
  class BoundAdd (line 23) | class BoundAdd(Bound):
    method __init__ (line 24) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 32) | def forward(self, x, y):
    method bound_backward (line 37) | def bound_backward(self, last_lA, last_uA, x, y, **kwargs):
    method bound_forward (line 49) | def bound_forward(self, dim_in, x, y):
    method interval_propagate (line 67) | def interval_propagate(self, x, y):
    method build_solver (line 71) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 108) | def build_gradient_node(self, grad_upstream):
  class BoundSub (line 124) | class BoundSub(Bound):
    method __init__ (line 125) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 130) | def forward(self, x, y):
    method bound_backward (line 135) | def bound_backward(self, last_lA, last_uA, x, y, **kwargs):
    method bound_forward (line 157) | def bound_forward(self, dim_in, x, y):
    method interval_propagate (line 176) | def interval_propagate(self, x, y):
    method build_solver (line 179) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 205) | def build_gradient_node(self, grad_upstream):
  class AddGrad (line 221) | class AddGrad(Module):
    method __init__ (line 222) | def __init__(self, input_shape, w=1.0):
    method forward (line 228) | def forward(self, grad_last):

FILE: auto_LiRPA/operators/base.py
  function not_implemented_op (line 35) | def not_implemented_op(node, func):
  class Interval (line 44) | class Interval(tuple):
    method __new__ (line 48) | def __new__(self, lb=None, ub=None, ptb=None):
    method __init__ (line 51) | def __init__(self, lb, ub, ptb=None):
    method __str__ (line 64) | def __str__(self):
    method __repr__ (line 67) | def __repr__(self):
    method make_interval (line 71) | def make_interval(lb, ub, other=None):
    method get_perturbation (line 79) | def get_perturbation(interval):
    method is_perturbed (line 98) | def is_perturbed(interval):
  class Bound (line 106) | class Bound(nn.Module):
    method __init__ (line 124) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method __repr__ (line 191) | def __repr__(self, attrs=None):
    method lower (line 202) | def lower(self):
    method lower (line 206) | def lower(self, value):
    method upper (line 216) | def upper(self):
    method upper (line 220) | def upper(self, value):
    method move_lower_and_upper_bounds_to_cache (line 229) | def move_lower_and_upper_bounds_to_cache(self):
    method delete_lower_and_upper_bounds (line 237) | def delete_lower_and_upper_bounds(self):
    method is_lower_bound_current (line 243) | def is_lower_bound_current(self):
    method is_upper_bound_current (line 246) | def is_upper_bound_current(self):
    method are_output_constraints_activated_for_layer (line 249) | def are_output_constraints_activated_for_layer(
    method init_gammas (line 271) | def init_gammas(self, num_constraints):
    method clip_gammas (line 296) | def clip_gammas(self):
    method is_input_perturbed (line 309) | def is_input_perturbed(self, i=0):
    method clear (line 313) | def clear(self):
    method input_name (line 318) | def input_name(self):
    method forward (line 321) | def forward(self, *x):
    method interval_propagate (line 333) | def interval_propagate(self, *v):
    method default_interval_propagate (line 354) | def default_interval_propagate(self, *v):
    method bound_forward (line 370) | def bound_forward(self, dim_in, *x):
    method bound_dynamic_forward (line 398) | def bound_dynamic_forward(self, *x, max_dim=None, offset=0):
    method bound_backward (line 401) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method broadcast_backward (line 421) | def broadcast_backward(self, A, x):
    method build_gradient_node (line 449) | def build_gradient_node(self, grad_upstream):
    method get_bias (line 468) | def get_bias(self, A, bias):
    method make_axis_non_negative (line 524) | def make_axis_non_negative(self, axis, shape='input'):
    method update_requires_input_bounds (line 552) | def update_requires_input_bounds(self):
    method clamp_interim_bounds (line 559) | def clamp_interim_bounds(self):
    method check_constraint_available (line 563) | def check_constraint_available(self, node, flag=False):
    method _ibp_constraint (line 571) | def _ibp_constraint(self, node: 'Bound', delete_bounds_after_use=False):
    method _check_weight_perturbation (line 614) | def _check_weight_perturbation(self):
    method non_deter_wrapper (line 626) | def non_deter_wrapper(self, op, *args, **kwargs):
    method non_deter_scatter_add (line 636) | def non_deter_scatter_add(self, *args, **kwargs):
    method non_deter_index_select (line 639) | def non_deter_index_select(self, *args, **kwargs):

FILE: auto_LiRPA/operators/bivariate.py
  class MulHelper (line 29) | class MulHelper:
    method __init__ (line 36) | def __init__(self):
    method interpolated_relaxation (line 40) | def interpolated_relaxation(x_l: Tensor, x_u: Tensor,
    method get_relaxation (line 78) | def get_relaxation(x_l: Tensor, x_u: Tensor, y_l: Tensor, y_u: Tensor,
    method get_forward_relaxation (line 95) | def get_forward_relaxation(x_l, x_u, y_l, y_u, opt_stage, alpha, start...
    method _get_gap (line 105) | def _get_gap(x, y, alpha, beta):
  class BoundMul (line 109) | class BoundMul(BoundOptimizableActivation):
    method __init__ (line 110) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 118) | def forward(self, x, y):
    method get_relaxation_opt (line 123) | def get_relaxation_opt(self, x_l, x_u, y_l, y_u):
    method _init_opt_parameters_impl (line 128) | def _init_opt_parameters_impl(self, size_spec, **kwargs):
    method _is_softmax (line 137) | def _is_softmax(self):
    method bound_relax (line 149) | def bound_relax(self, x, y, init=False, dim_opt=None):
    method _multiply_by_const (line 194) | def _multiply_by_const(x, const):
    method bound_backward_constant (line 216) | def bound_backward_constant(self, last_lA, last_uA, x, y, op=None,
    method bound_backward (line 237) | def bound_backward(self, last_lA, last_uA, x, y, start_node=None, **kw...
    method bound_backward_both_perturbed (line 247) | def bound_backward_both_perturbed(self, last_lA, last_uA, x, y,
    method bound_forward (line 327) | def bound_forward(self, dim_in, x, y):
    method bound_forward_constant (line 337) | def bound_forward_constant(self, x, y, batched_constant):
    method bound_forward_both_perturbed (line 354) | def bound_forward_both_perturbed(self, dim_in, x, y):
    method interval_propagate_constant (line 379) | def interval_propagate_constant(x, y, op=lambda x, const: x * const):
    method interval_propagate (line 390) | def interval_propagate(self, x, y):
    method interval_propagate_both_perturbed (line 406) | def interval_propagate_both_perturbed(*v):
    method build_solver (line 425) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method update_requires_input_bounds (line 433) | def update_requires_input_bounds(self):
    method build_gradient_node (line 449) | def build_gradient_node(self, grad_upstream):
  class MulGrad (line 458) | class MulGrad(Module):
    method __init__ (line 459) | def __init__(self, input_shape):
    method forward (line 464) | def forward(self, grad_last, y):
  class BoundDiv (line 473) | class BoundDiv(Bound):
    method forward (line 475) | def forward(self, x, y):

FILE: auto_LiRPA/operators/clampmult.py
  class ClampedMultiplication (line 28) | class ClampedMultiplication(torch.autograd.Function):
    method clamp_mutiply_forward (line 32) | def clamp_mutiply_forward(A: Tensor, d_pos: Tensor, d_neg: Tensor,
    method clamp_mutiply_backward (line 68) | def clamp_mutiply_backward(A: Tensor, d_pos: Tensor, d_neg: Tensor,
    method forward (line 130) | def forward(ctx, A, d_pos, d_neg, b_pos, b_neg, patches_mode, reduce_b...
    method backward (line 140) | def backward(ctx, grad_output_A, grad_output_bias):
  function multiply_by_A_signs (line 148) | def multiply_by_A_signs(A, d_pos, d_neg, b_pos, b_neg, contiguous='auto',

FILE: auto_LiRPA/operators/constant.py
  class BoundConstant (line 21) | class BoundConstant(Bound):
    method __init__ (line 22) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method __repr__ (line 28) | def __repr__(self):
    method forward (line 34) | def forward(self):
    method bound_backward (line 37) | def bound_backward(self, last_lA, last_uA, **kwargs):
    method bound_forward (line 57) | def bound_forward(self, dim_in):
    method build_solver (line 62) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
  class BoundPrimConstant (line 66) | class BoundPrimConstant(Bound):
    method forward (line 67) | def forward(self):
  class BoundConstantOfShape (line 71) | class BoundConstantOfShape(Bound):
    method __init__ (line 72) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 77) | def forward(self, x):
    method bound_backward (line 82) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method bound_forward (line 99) | def bound_forward(self, dim_in, x):
    method interval_propagate (line 105) | def interval_propagate(self, *v):
    method build_solver (line 110) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
  class BoundRange (line 114) | class BoundRange(Bound):
    method __init__ (line 115) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 119) | def forward(self, start, end, step):
  class BoundATenDiag (line 126) | class BoundATenDiag(Bound):
    method __init__ (line 127) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 131) | def forward(self, x, diagonal=0):
    method interval_propagate (line 134) | def interval_propagate(self, *v):
  class BoundATenDiagonal (line 138) | class BoundATenDiagonal(Bound):
    method __init__ (line 139) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 143) | def forward(self, x, offset=0, dim1=0, dim2=1):
    method interval_propagate (line 146) | def interval_propagate(self, *v):
    method bound_backward (line 150) | def bound_backward(self, last_lA, last_uA, *args, **kwargs):

FILE: auto_LiRPA/operators/convex_concave.py
  class BoundLog (line 24) | class BoundLog(BoundActivation):
    method __init__ (line 26) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 30) | def forward(self, x):
    method bound_relax (line 36) | def bound_relax(self, x, init=False):
    method interval_propagate (line 47) | def interval_propagate(self, *v):
    method bound_backward (line 56) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
  class BoundSqrt (line 66) | class BoundSqrt(BoundOptimizableActivation):
    method __init__ (line 67) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 73) | def forward(self, x):
    method bound_relax (line 76) | def bound_relax(self, x, init=False, dim_opt=None):
    method bound_backward (line 94) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method clamp_interim_bounds (line 105) | def clamp_interim_bounds(self):
    method _init_opt_parameters_impl (line 110) | def _init_opt_parameters_impl(self, size_spec, **kwargs):
  class BoundReciprocal (line 118) | class BoundReciprocal(BoundOptimizableActivation):
    method __init__ (line 120) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 125) | def forward(self, x):
    method interval_propagate (line 128) | def interval_propagate(self, *v):
    method bound_relax (line 134) | def bound_relax(self, x, init=False, dim_opt=None):
    method bound_backward (line 162) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method _init_opt_parameters_impl (line 170) | def _init_opt_parameters_impl(self, size_spec, **kwargs):
    method build_gradient_node (line 177) | def build_gradient_node(self, grad_upstream):
  class ReciprocalGrad (line 181) | class ReciprocalGrad(Module):
    method __init__ (line 182) | def __init__(self):
    method forward (line 185) | def forward(self, g, x):
  class BoundExp (line 190) | class BoundExp(BoundOptimizableActivation):
    method __init__ (line 191) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 198) | def forward(self, x):
    method interval_propagate (line 204) | def interval_propagate(self, *v):
    method bound_forward (line 215) | def bound_forward(self, dim_in, x):
    method bound_backward (line 230) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method _check_nan (line 274) | def _check_nan(self, A, bias, last_A, const_bound):
    method bound_relax (line 298) | def bound_relax(self, x, init=False, dim_opt=None):
    method _init_opt_parameters_impl (line 315) | def _init_opt_parameters_impl(self, size_spec, **kwargs):
    method build_gradient_node (line 322) | def build_gradient_node(self, grad_upstream):
  class ExpGrad (line 328) | class ExpGrad(Module):
    method __init__ (line 329) | def __init__(self):
    method forward (line 332) | def forward(self, g, preact):

FILE: auto_LiRPA/operators/convolution.py
  class BoundConv (line 27) | class BoundConv(Bound):
    method __init__ (line 28) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 62) | def forward(self, *x):
    method bound_backward (line 70) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method build_solver (line 212) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method interval_propagate (line 349) | def interval_propagate(self, *v, C=None):
    method bound_dynamic_forward (line 394) | def bound_dynamic_forward(self, *x, max_dim=None, offset=0):
    method bound_forward (line 426) | def bound_forward(self, dim_in, *x):
    method build_gradient_node (line 462) | def build_gradient_node(self, grad_upstream):
    method update_requires_input_bounds (line 468) | def update_requires_input_bounds(self):
  class BoundConvTranspose (line 472) | class BoundConvTranspose(Bound):
    method __init__ (line 473) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 496) | def forward(self, *x):
    method bound_backward (line 503) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method interval_propagate (line 617) | def interval_propagate(self, *v, C=None):
    method bound_forward (line 650) | def bound_forward(self, dim_in, *x):
  class BoundPad (line 687) | class BoundPad(Bound):
    method __init__ (line 688) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 697) | def forward(self, x, pad, value=0.0):
    method interval_propagate (line 707) | def interval_propagate(self, *v):
    method bound_backward (line 711) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method build_solver (line 731) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
  class Conv2dGrad (line 772) | class Conv2dGrad(Module):
    method __init__ (line 773) | def __init__(self, fw_module, weight, stride, padding, dilation, groups):
    method forward (line 787) | def forward(self, grad_last):
  class Conv2dGradOp (line 802) | class Conv2dGradOp(Function):
    method symbolic (line 804) | def symbolic(g, x, w, stride, padding, dilation, groups,
    method forward (line 813) | def forward(
  class BoundConv2dGrad (line 826) | class BoundConv2dGrad(Bound):
    method __init__ (line 827) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 841) | def forward(self, *x):
    method bound_backward (line 848) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method interval_propagate (line 956) | def interval_propagate(self, *v, C=None):

FILE: auto_LiRPA/operators/cut_ops.py
  class CutModule (line 22) | class CutModule():
    method __init__ (line 24) | def __init__(self, relu_nodes=[], general_beta=None, x_coeffs=None,
    method use_patches (line 44) | def use_patches(self, start_node):
    method select_active_general_beta (line 49) | def select_active_general_beta(self, start_node, unstable_idx=None):
    method general_beta_coeffs_mm (line 82) | def general_beta_coeffs_mm(self, unstable_spec_beta, coeffs, A, curren...
    method general_beta_coeffs_addmm_to_A (line 108) | def general_beta_coeffs_addmm_to_A(self, lA, uA, general_beta, coeffs,...
    method patch_trick (line 134) | def patch_trick(self, start_node, layer_name, A, current_layer_shape):
    method relu_cut (line 166) | def relu_cut(self, start_node, layer_name, last_lA, last_uA, current_l...
    method pre_cut (line 185) | def pre_cut(self, start_node, layer_name, lA, uA, current_layer_shape,...
    method jit_arelu_lA (line 204) | def jit_arelu_lA(last_lA, lower, upper, beta_mm_coeffs, unstable_or_cu...
    method jit_arelu_lbias (line 230) | def jit_arelu_lbias(unstable_or_cut_index, nu_hat_pos, beta_mm_coeffs,...
    method jit_arelu_uA (line 258) | def jit_arelu_uA(last_uA, lower, upper, beta_mm_coeffs, unstable_or_cu...
    method jit_arelu_ubias (line 277) | def jit_arelu_ubias(unstable_or_cut_index, nu_hat_pos, beta_mm_coeffs,...
    method arelu_cut (line 298) | def arelu_cut(self, start_node, layer_name, last_lA, last_uA, lower_d,...
    method input_cut (line 493) | def input_cut(self, start_node, lA, uA, current_layer_shape, unstable_...
    method bias_cut (line 513) | def bias_cut(self, start_node, lb, ub, unstable_idx=None, batch_mask=N...
  function _bound_oneside (line 533) | def _bound_oneside(last_A, d_pos, d_neg, b_pos, b_neg, start_node, patch...
  function _maybe_unfold (line 576) | def _maybe_unfold(d_tensor, last_A):

FILE: auto_LiRPA/operators/dropout.py
  class BoundDropout (line 19) | class BoundDropout(Bound):
    method __init__ (line 20) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method clear (line 30) | def clear(self):
    method forward (line 33) | def forward(self, *inputs):
    method _check_forward (line 50) | def _check_forward(self):
    method bound_backward (line 56) | def bound_backward(self, last_lA, last_uA, *args, **kwargs):
    method bound_forward (line 69) | def bound_forward(self, dim_in, x, *args):
    method interval_propagate (line 79) | def interval_propagate(self, *v):

FILE: auto_LiRPA/operators/dtype.py
  class BoundCast (line 20) | class BoundCast(Bound):
    method __init__ (line 21) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 36) | def forward(self, x):
    method bound_backward (line 40) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method bound_forward (line 51) | def bound_forward(self, dim_in, x):
    method build_solver (line 56) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...

FILE: auto_LiRPA/operators/gelu.py
  class BoundGelu (line 26) | class BoundGelu(BoundTanh):
    method __init__ (line 29) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method _init_masks (line 43) | def _init_masks(self, x):
    method precompute_relaxation (line 67) | def precompute_relaxation(self, func, dfunc, x_limit=1000):
    method opt_init (line 186) | def opt_init(self):
    method _init_opt_parameters_impl (line 194) | def _init_opt_parameters_impl(self, size_spec, name_start):
    method forward (line 207) | def forward(self, x):
    method bound_relax_impl (line 210) | def bound_relax_impl(self, x, func, dfunc):
    method bound_relax (line 385) | def bound_relax(self, x, init=False, dim_opt=None):
    method interval_propagate (line 390) | def interval_propagate(self, *v):
  class GELUOp (line 398) | class GELUOp(torch.autograd.Function):
    method symbolic (line 403) | def symbolic(g, x):
    method forward (line 407) | def forward(ctx, x):
    method backward (line 412) | def backward(ctx, grad_output):
  class GELU (line 420) | class GELU(nn.Module):
    method forward (line 421) | def forward(self, x):

FILE: auto_LiRPA/operators/indexing.py
  class BoundGather (line 22) | class BoundGather(Bound):
    method __init__ (line 23) | def __init__(self, attr, x, output_index, options):
    method forward (line 27) | def forward(self, x, indices):
    method bound_backward (line 47) | def bound_backward(self, last_lA, last_uA, *args, **kwargs):
    method bound_forward (line 102) | def bound_forward(self, dim_in, x, indices):
    method interval_propagate (line 122) | def interval_propagate(self, *v):
    method build_solver (line 126) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 129) | def build_gradient_node(self, grad_upstream):
  class GatherGrad (line 133) | class GatherGrad(Module):
    method __init__ (line 134) | def __init__(self, axis, indices, input_shape):
    method forward (line 140) | def forward(self, grad_last):
  class BoundGatherElements (line 216) | class BoundGatherElements(Bound):
    method __init__ (line 217) | def __init__(self, attr, input, output_index, options):
    method forward (line 221) | def forward(self, x, index):
    method bound_backward (line 225) | def bound_backward(self, last_lA, last_uA, x, index, **kwargs):
    method interval_propagate (line 243) | def interval_propagate(self, *v):
    method bound_forward (line 248) | def bound_forward(self, dim_in, x, index):
    method _get_dim (line 257) | def _get_dim(self):

FILE: auto_LiRPA/operators/jacobian.py
  class JacobianOP (line 23) | class JacobianOP(torch.autograd.Function):
    method symbolic (line 25) | def symbolic(g, output, input):
    method forward (line 29) | def forward(ctx, output, input):
  class BoundJacobianOP (line 36) | class BoundJacobianOP(Bound):
    method __init__ (line 37) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 40) | def forward(self, output, input):
  class BoundJacobianInit (line 44) | class BoundJacobianInit(Bound):
    method __init__ (line 45) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 49) | def forward(self, x):
  class GradNorm (line 58) | class GradNorm(Module):
    method __init__ (line 59) | def __init__(self, norm=1):
    method forward (line 63) | def forward(self, grad):

FILE: auto_LiRPA/operators/leaf.py
  class BoundInput (line 25) | class BoundInput(Bound):
    method __init__ (line 26) | def __init__(self, ori_name, value, perturbation=None, input_index=Non...
    method __setattr__ (line 35) | def __setattr__(self, key, value):
    method forward (line 44) | def forward(self):
    method bound_forward (line 47) | def bound_forward(self, dim_in):
    method bound_backward (line 50) | def bound_backward(self, last_lA, last_uA, **kwargs):
    method interval_propagate (line 54) | def interval_propagate(self, *v):
  class BoundParams (line 58) | class BoundParams(BoundInput):
    method __init__ (line 59) | def __init__(self, ori_name, value, perturbation=None, options=None, a...
    method register_parameter (line 67) | def register_parameter(self, name, param):
    method init (line 75) | def init(self, initializing=False):
    method forward (line 78) | def forward(self):
  class BoundBuffers (line 84) | class BoundBuffers(BoundInput):
    method __init__ (line 85) | def __init__(self, ori_name, value, perturbation=None, options=None, a...
    method forward (line 97) | def forward(self):

FILE: auto_LiRPA/operators/linear.py
  class BoundLinear (line 31) | class BoundLinear(BoundOptimizableActivation):
    method __init__ (line 32) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method _preprocess (line 66) | def _preprocess(self, a, b, c=None):
    method init_opt_parameters (line 81) | def init_opt_parameters(self, start_nodes):
    method forward (line 113) | def forward(self, x, w, b=None):
    method onehot_mult (line 122) | def onehot_mult(self, weight, bias, C, batch_size):
    method bound_backward (line 170) | def bound_backward(self, last_lA, last_uA, *x, start_node=None,
    method _reshape (line 481) | def _reshape(self, x_l, x_u, y_l, y_u):
    method propagate_A_xy (line 504) | def propagate_A_xy(last_A: Tensor, alpha_pos: Tensor, alpha_neg: Tensor,
    method bound_backward_with_weight (line 524) | def bound_backward_with_weight(self, last_lA, last_uA, input_lb, input...
    method _propagate_Linf (line 600) | def _propagate_Linf(x, w):
    method interval_propagate (line 613) | def interval_propagate(self, *v, C=None, w=None):
    method interval_propagate_with_weight (line 692) | def interval_propagate_with_weight(self, *v):
    method bound_forward_mul (line 748) | def bound_forward_mul(x_lw: Tensor, x_lb: Tensor, x_uw: Tensor, x_ub: ...
    method bound_dynamic_forward (line 765) | def bound_dynamic_forward(self, x, w=None, b=None, C=None, max_dim=Non...
    method bound_forward (line 791) | def bound_forward(self, dim_in, x, w=None, b=None, C=None, weight_has_...
    method bound_forward_with_weight (line 840) | def bound_forward_with_weight(self, dim_in, x, y):
    method build_solver (line 867) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 938) | def build_gradient_node(self, grad_upstream):
    method update_requires_input_bounds (line 954) | def update_requires_input_bounds(self):
  class BoundMatMul (line 958) | class BoundMatMul(BoundLinear):
    method __init__ (line 960) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 966) | def forward(self, x, y):
    method interval_propagate (line 971) | def interval_propagate(self, *v, C=None):
    method bound_backward (line 975) | def bound_backward(self, last_lA, last_uA, *x, start_node=None, **kwar...
    method bound_forward (line 1003) | def bound_forward(self, dim_in, x, y):
    method update_requires_input_bounds (line 1022) | def update_requires_input_bounds(self):
  class BoundNeg (line 1036) | class BoundNeg(Bound):
    method __init__ (line 1037) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 1041) | def forward(self, x):
    method bound_backward (line 1044) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method bound_forward (line 1066) | def bound_forward(self, dim_in, x):
    method interval_propagate (line 1069) | def interval_propagate(self, *v):
    method build_gradient_node (line 1072) | def build_gradient_node(self, grad_upstream):
  class NegGrad (line 1076) | class NegGrad(Module):
    method forward (line 1077) | def forward(self, grad_last):
  class BoundCumSum (line 1081) | class BoundCumSum(Bound):
    method __init__ (line 1082) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 1086) | def forward(self, x, axis):
  class BoundIdentity (line 1090) | class BoundIdentity(Bound):
    method __init__ (line 1091) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 1095) | def forward(self, x):
    method bound_backward (line 1098) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method bound_forward (line 1101) | def bound_forward(self, dim_in, x):
  class LinearGrad (line 1105) | class LinearGrad(Module):
    method __init__ (line 1106) | def __init__(self, weight):
    method forward (line 1110) | def forward(self, grad_last):
  class MatMulGrad (line 1115) | class MatMulGrad(Module):
    method forward (line 1116) | def forward(self, grad_last, x):

FILE: auto_LiRPA/operators/logical.py
  class BoundWhere (line 21) | class BoundWhere(Bound):
    method forward (line 22) | def forward(self, condition, x, y):
    method interval_propagate (line 25) | def interval_propagate(self, *v):
    method bound_backward (line 30) | def bound_backward(self, last_lA, last_uA, condition, x, y, **kwargs):
  class BoundNot (line 48) | class BoundNot(Bound):
    method forward (line 49) | def forward(self, x):
  class BoundEqual (line 53) | class BoundEqual(Bound):
    method forward (line 54) | def forward(self, x, y):

FILE: auto_LiRPA/operators/minmax.py
  class BoundMinMax (line 23) | class BoundMinMax(BoundOptimizableActivation):
    method __init__ (line 24) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method _init_opt_parameters_impl (line 30) | def _init_opt_parameters_impl(self, size_spec, name_start):
    method clip_alpha (line 37) | def clip_alpha(self):
    method forward (line 92) | def forward(self, x, y):
    method _backward_relaxation (line 100) | def _backward_relaxation(self, x, y, start_node=None):
    method bound_backward (line 221) | def bound_backward(self, last_lA, last_uA, x=None, y=None, start_shape...
    method interval_propagate (line 277) | def interval_propagate(self, *v):
  class BoundMax (line 283) | class BoundMax(BoundMinMax):
    method __init__ (line 284) | def __init__(self, *args, **kwargs):
  class BoundMin (line 289) | class BoundMin(BoundMinMax):
    method __init__ (line 290) | def __init__(self, *args, **kwargs):

FILE: auto_LiRPA/operators/normalization.py
  class BoundBatchNormalization (line 29) | class BoundBatchNormalization(Bound):
    method __init__ (line 30) | def __init__(self, attr, inputs, output_index, options, training):
    method _check_unused_mean_or_var (line 48) | def _check_unused_mean_or_var(self):
    method forward (line 55) | def forward(self, x, w, b, m, v):
    method bound_forward (line 79) | def bound_forward(self, dim_in, *x):
    method bound_backward (line 112) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method interval_propagate (line 204) | def interval_propagate(self, *v):
    method build_solver (line 257) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method update_requires_input_bounds (line 299) | def update_requires_input_bounds(self):
  class LayerNormImpl (line 303) | class LayerNormImpl(nn.Module):
    method __init__ (line 304) | def __init__(self, axis, epsilon):
    method forward (line 309) | def forward(self, x, scale, bias):
  class BoundLayerNormalization (line 322) | class BoundLayerNormalization(Bound):
    method __init__ (line 323) | def __init__(self, attr, inputs, output_index, options):
    method forward (line 328) | def forward(self, x, scale, bias):

FILE: auto_LiRPA/operators/pooling.py
  class BoundMaxPool (line 25) | class BoundMaxPool(BoundOptimizableActivation):
    method __init__ (line 27) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 41) | def forward(self, x):
    method project_simplex (line 46) | def project_simplex(self, patches):
    method _init_opt_parameters_impl (line 60) | def _init_opt_parameters_impl(self, size_spec, name_start):
    method jit_mutiply (line 75) | def jit_mutiply(Apos, Aneg, pos, neg):
    method bound_backward (line 78) | def bound_backward(self, last_lA, last_uA, x, start_node=None,
    method bound_forward (line 339) | def bound_forward(self, dim_in, x):
    method bound_relax (line 365) | def bound_relax(self, x, init=False, dim_opt=None):
    method dump_alpha (line 416) | def dump_alpha(self, device=None, dtype=None, non_blocking=False):
    method restore_alpha (line 421) | def restore_alpha(self, alpha, device=None, dtype=None, non_blocking=F...
    method drop_unused_alpha (line 425) | def drop_unused_alpha(self, keep_nodes):
    method build_solver (line 431) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
  class BoundGlobalAveragePool (line 479) | class BoundGlobalAveragePool(Bound):
    method __init__ (line 480) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 483) | def forward(self, x):
    method bound_backward (line 487) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method interval_propagate (line 495) | def interval_propagate(self, *v):
  class BoundAveragePool (line 502) | class BoundAveragePool(Bound):
    method __init__ (line 503) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 525) | def forward(self, x):
    method bound_backward (line 529) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method build_solver (line 663) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...

FILE: auto_LiRPA/operators/reduce.py
  class BoundReduce (line 22) | class BoundReduce(Bound):
    method __init__ (line 23) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method _parse_input_and_axis (line 29) | def _parse_input_and_axis(self, *x):
    method _return_bound_backward (line 36) | def _return_bound_backward(self, lA, uA):
  class BoundReduceMax (line 40) | class BoundReduceMax(BoundReduce):
    method __init__ (line 41) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method _parse_input_and_axis (line 48) | def _parse_input_and_axis(self, *x):
    method forward (line 56) | def forward(self, *x):
    method bound_backward (line 62) | def bound_backward(self, last_lA, last_uA, *args, **kwargs):
    method build_gradient_node (line 86) | def build_gradient_node(self, grad_upstream):
  class ReduceMaxGrad (line 96) | class ReduceMaxGrad(Module):
    method __init__ (line 97) | def __init__(self, axis, keepdim, input_shape, indices):
    method forward (line 104) | def forward(self, grad_last):
  class BoundReduceMin (line 154) | class BoundReduceMin(BoundReduceMax):
    method forward (line 155) | def forward(self, *x):
  class BoundReduceMean (line 162) | class BoundReduceMean(BoundReduce):
    method forward (line 163) | def forward(self, *x):
    method bound_backward (line 167) | def bound_backward(self, last_lA, last_uA, *args, **kwargs):
    method bound_forward (line 186) | def bound_forward(self, dim_in, x, *args):
  class BoundReduceSum (line 199) | class BoundReduceSum(BoundReduce):
    method forward (line 200) | def forward(self, *x):
    method bound_backward (line 207) | def bound_backward(self, last_lA, last_uA, x, *args, **kwargs):
    method bound_forward (line 229) | def bound_forward(self, dim_in, x, *args):
    method build_gradient_node (line 241) | def build_gradient_node(self, grad_upstream):
  class ReduceSumGrad (line 246) | class ReduceSumGrad(Module):
    method __init__ (line 247) | def __init__(self, axis, keepdim, input_shape):
    method forward (line 253) | def forward(self, grad_last):

FILE: auto_LiRPA/operators/relu.py
  class BoundTwoPieceLinear (line 31) | class BoundTwoPieceLinear(BoundOptimizableActivation):
    method __init__ (line 32) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method init_opt_parameters (line 52) | def init_opt_parameters(self, start_nodes):
    method select_alpha_by_idx (line 164) | def select_alpha_by_idx(self, last_lA, last_uA, unstable_idx, start_no...
    method reconstruct_full_alpha (line 232) | def reconstruct_full_alpha(self, sparse_alpha, full_alpha_shape, alpha...
    method bound_backward (line 247) | def bound_backward(self, last_lA, last_uA, x=None, start_node=None,
    method _transfer_alpha_lookup_idx (line 361) | def _transfer_alpha_lookup_idx(self, alpha_lookup_idx, device=None, dt...
    method _transfer_alpha_indices (line 368) | def _transfer_alpha_indices(self, alpha_indices, device=None, dtype=No...
    method dump_alpha (line 374) | def dump_alpha(self, device=None, dtype=None, non_blocking=False):
    method restore_alpha (line 382) | def restore_alpha(self, alpha, device=None, dtype=None, non_blocking=F...
    method drop_unused_alpha (line 389) | def drop_unused_alpha(self, keep_nodes):
  class BoundRelu (line 404) | class BoundRelu(BoundTwoPieceLinear):
    method __init__ (line 405) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method get_unstable_idx (line 413) | def get_unstable_idx(self):
    method clip_alpha (line 417) | def clip_alpha(self):
    method forward (line 421) | def forward(self, x):
    method _relu_lower_bound_init (line 430) | def _relu_lower_bound_init(self, upper_k):
    method _relu_upper_opt_same_slope (line 456) | def _relu_upper_opt_same_slope(self, lb_lower_d, ub_lower_d, upper_d, ...
    method _forward_relaxation (line 493) | def _forward_relaxation(self, x):
    method bound_dynamic_forward (line 520) | def bound_dynamic_forward(self, x, max_dim=None, offset=0):
    method bound_forward (line 571) | def bound_forward(self, dim_in, x):
    method _relu_upper_bound (line 584) | def _relu_upper_bound(lb, ub, leaky_alpha: float):
    method _relu_mask_alpha (line 598) | def _relu_mask_alpha(lower, upper, lb_lower_d : Optional[Tensor],
    method _backward_relaxation (line 622) | def _backward_relaxation(self, last_lA, last_uA, x, start_node, unstab...
    method interval_propagate (line 708) | def interval_propagate(self, *v):
    method build_solver (line 712) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 782) | def build_gradient_node(self, grad_upstream):
    method get_split_mask (line 791) | def get_split_mask(self, lower, upper, input_index):
    method get_unstable_mask (line 796) | def get_unstable_mask(self, lower, upper):
    method compute_bound_improvement_heuristics (line 804) | def compute_bound_improvement_heuristics(self, lower, upper):
  class BoundLeakyRelu (line 812) | class BoundLeakyRelu(BoundRelu):
  class BoundSign (line 816) | class BoundSign(BoundActivation):
    method __init__ (line 817) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 821) | def forward(self, x):
    method bound_relax (line 824) | def bound_relax(self, x, init=False):
  class SignMergeFunction_loose (line 856) | class SignMergeFunction_loose(torch.autograd.Function):
    method forward (line 860) | def forward(ctx, input):
    method backward (line 866) | def backward(ctx, grad_output):
  class SignMergeFunction_tight (line 874) | class SignMergeFunction_tight(torch.autograd.Function):
    method forward (line 878) | def forward(ctx, input):
    method backward (line 884) | def backward(ctx, grad_output):
  class BoundSignMerge (line 893) | class BoundSignMerge(BoundTwoPieceLinear):
    method __init__ (line 894) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method get_unstable_idx (line 901) | def get_unstable_idx(self):
    method forward (line 905) | def forward(self, x):
    method _mask_alpha (line 909) | def _mask_alpha(self, lower, upper, lb_lower_d, ub_lower_d, lb_upper_d...
    method _backward_relaxation (line 925) | def _backward_relaxation(self, last_lA, last_uA, x, start_node, unstab...
    method build_solver (line 989) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
  function relu_grad (line 1042) | def relu_grad(preact):
  class ReLUGradOp (line 1046) | class ReLUGradOp(Function):
    method symbolic (line 1052) | def symbolic(_, g, g_relu, g_relu_rev, preact):
    method forward (line 1056) | def forward(ctx, g, g_relu, g_relu_rev, preact):
  class ReLUGrad (line 1060) | class ReLUGrad(Module):
    method forward (line 1061) | def forward(self, g, preact):
  function _maybe_unfold (line 1067) | def _maybe_unfold(d_tensor, last_A):
  class BoundReluGrad (line 1101) | class BoundReluGrad(BoundActivation):
    method __init__ (line 1102) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method relu_grad (line 1108) | def relu_grad(preact):
    method forward (line 1111) | def forward(self, g, g_relu, g_relu_rev, preact):
    method interval_propagate (line 1116) | def interval_propagate(self, *v):
    method bound_backward (line 1128) | def bound_backward(self, last_lA, last_uA, g, g_relu, g_relu_rev, preact,

FILE: auto_LiRPA/operators/reshape.py
  class BoundReshape (line 24) | class BoundReshape(Bound):
    method __init__ (line 25) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 30) | def forward(self, x, shape):
    method bound_backward (line 41) | def bound_backward(self, last_lA, last_uA, x, shape, **kwargs):
    method bound_forward (line 74) | def bound_forward(self, dim_in, x, shape):
    method bound_dynamic_forward (line 82) | def bound_dynamic_forward(self, x, shape, max_dim=None, offset=0):
    method interval_propagate (line 87) | def interval_propagate(self, *v):
    method build_solver (line 92) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 100) | def build_gradient_node(self, grad_upstream):
  class BoundUnsqueeze (line 106) | class BoundUnsqueeze(Bound):
    method __init__ (line 107) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 117) | def forward(self, *x):
    method bound_backward (line 126) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method bound_forward (line 148) | def bound_forward(self, dim_in, *x):
    method build_solver (line 160) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 163) | def build_gradient_node(self, grad_upstream):
  class UnsqueezeGrad (line 171) | class UnsqueezeGrad(Module):
    method __init__ (line 172) | def __init__(self, axes):
    method forward (line 176) | def forward(self, grad_last):
  class BoundExpand (line 180) | class BoundExpand(Bound):
    method __init__ (line 181) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 185) | def forward(self, x, y):
    method bound_backward (line 197) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method bound_forward (line 216) | def bound_forward(self, dim_in, *x):
    method build_gradient_node (line 225) | def build_gradient_node(self, grad_upstream):
  class ExpandGrad (line 232) | class ExpandGrad(Module):
    method __init__ (line 236) | def __init__(self, shape):
    method forward (line 240) | def forward(self, grad_last):
  class BoundSqueeze (line 244) | class BoundSqueeze(Bound):
    method __init__ (line 245) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 255) | def forward(self, *x):
    method bound_backward (line 263) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method bound_forward (line 274) | def bound_forward(self, dim_in, *x):
    method build_solver (line 287) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
  class BoundFlatten (line 291) | class BoundFlatten(Bound):
    method __init__ (line 292) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 297) | def forward(self, x):
    method bound_backward (line 300) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method bound_dynamic_forward (line 307) | def bound_dynamic_forward(self, x, max_dim=None, offset=0):
    method bound_forward (line 312) | def bound_forward(self, dim_in, x):
    method build_solver (line 322) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 327) | def build_gradient_node(self, grad_upstream):
  class BoundATenUnflatten (line 333) | class BoundATenUnflatten(BoundReshape):
    method __init__ (line 334) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 337) | def forward(self, x, dim, sizes):
    method bound_backward (line 344) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method bound_forward (line 350) | def bound_forward(self, dim_in, *x):
    method bound_dynamic_forward (line 353) | def bound_dynamic_forward(self, *x, max_dim=None, offset=0):
    method interval_propagate (line 356) | def interval_propagate(self, x, dim, sizes):
    method build_solver (line 361) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
  class ReshapeGrad (line 366) | class ReshapeGrad(Module):
    method forward (line 367) | def forward(self, grad_last, inp):
  class BoundTranspose (line 374) | class BoundTranspose(Bound):
    method __init__ (line 375) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 385) | def forward(self, x):
    method bound_backward (line 388) | def bound_backward(self, last_lA, last_uA, x, **kwargs):
    method bound_forward (line 396) | def bound_forward(self, dim_in, x):
    method build_solver (line 407) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 410) | def build_gradient_node(self, grad_upstream):
  class TransposeGrad (line 416) | class TransposeGrad(Module):
    method __init__ (line 417) | def __init__(self, perm_inv):
    method forward (line 421) | def forward(self, grad_last):

FILE: auto_LiRPA/operators/resize.py
  class BoundResize (line 26) | class BoundResize(Bound):
    method __init__ (line 27) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 34) | def forward(self, x, size=None, scale_factor=None):
    method interval_propagate (line 49) | def interval_propagate(self, *v):
    method bound_forward (line 53) | def bound_forward(self, dim_in, *inp):
    method bound_backward (line 67) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):

FILE: auto_LiRPA/operators/rnn.py
  class BoundRNN (line 21) | class BoundRNN(Bound):
    method __init__ (line 22) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 32) | def forward(self, x, weight_input, weight_recurrent, bias, sequence_le...

FILE: auto_LiRPA/operators/s_shaped.py
  class BoundSShaped (line 25) | class BoundSShaped(BoundOptimizableActivation):
    method __init__ (line 30) | def __init__(self, attr=None, inputs=None, output_index=0, options=Non...
    method opt_init (line 72) | def opt_init(self):
    method branch_input_domain (line 77) | def branch_input_domain(self, lb, ub):
    method _init_opt_parameters_impl (line 85) | def _init_opt_parameters_impl(self, size_spec, name_start, num_params=...
    method precompute_relaxation (line 99) | def precompute_relaxation(self, func, dfunc, x_limit=500):
    method precompute_dfunc_values (line 167) | def precompute_dfunc_values(self, func, dfunc, x_limit=500):
    method forward (line 174) | def forward(self, x):
    method retrieve_from_precompute (line 177) | def retrieve_from_precompute(self, precomputed_d, input_bound, default...
    method generate_d_lower_upper (line 204) | def generate_d_lower_upper(self, lower, upper):
    method retrieve_d_from_k (line 216) | def retrieve_d_from_k(self, k, func):
    method bound_relax_impl_same_slope (line 232) | def bound_relax_impl_same_slope(self, x, func, dfunc):
    method bound_relax_impl (line 279) | def bound_relax_impl(self, x, func, dfunc):
    method bound_relax_branch (line 488) | def bound_relax_branch(self, lb, ub):
    method bound_relax (line 494) | def bound_relax(self, x, init=False, dim_opt=None):
    method get_split_mask (line 510) | def get_split_mask(self, lower, upper, input_index):
  class BoundPow (line 518) | class BoundPow(BoundSShaped):
    method __init__ (line 519) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method generate_d_lower_upper (line 535) | def generate_d_lower_upper(self, lower, upper):
    method branch_input_domain (line 549) | def branch_input_domain(self, lb, ub):
    method precompute_relaxation (line 581) | def precompute_relaxation(self, func, dfunc, x_limit = 500):
    method forward (line 650) | def forward(self, x, y):
    method bound_backward (line 653) | def bound_backward(self, last_lA, last_uA, x, y, start_node=None,
    method bound_forward (line 671) | def bound_forward(self, dim_in, x, y):
    method bound_relax_branch (line 682) | def bound_relax_branch(self, lb, ub):
    method bound_relax (line 702) | def bound_relax(self, x, init=False, dim_opt=None):
    method interval_propagate (line 711) | def interval_propagate(self, *v):
    method clamp_interim_bounds (line 724) | def clamp_interim_bounds(self):
  function dtanh (line 731) | def dtanh(x):
  function dsigmoid (line 734) | def dsigmoid(x):
  function darctan (line 737) | def darctan(x):
  function d2tanh (line 740) | def d2tanh(x):
  function d2sigmoid (line 743) | def d2sigmoid(x):
  class BoundTanh (line 747) | class BoundTanh(BoundSShaped):
    method __init__ (line 752) | def __init__(self, attr=None, inputs=None, output_index=0, options=None,
    method _init_opt_parameters_impl (line 757) | def _init_opt_parameters_impl(self, size_spec, name_start):
    method build_gradient_node (line 761) | def build_gradient_node(self, grad_upstream):
  class TanhGradOp (line 768) | class TanhGradOp(Function):
    method symbolic (line 770) | def symbolic(_, preact):
    method forward (line 774) | def forward(ctx, preact):
  class TanhGrad (line 778) | class TanhGrad(Module):
    method forward (line 779) | def forward(self, g, preact):
  class BoundTanhGrad (line 783) | class BoundTanhGrad(BoundOptimizableActivation):
    method __init__ (line 784) | def __init__(self, attr=None, inputs=None, output_index=0, options=None,
    method forward (line 795) | def forward(self, x):
    method interval_propagate (line 798) | def interval_propagate(self, *v):
    method bound_relax (line 808) | def bound_relax(self, x, init=False, dim_opt=None):
    method precompute_relaxation (line 813) | def precompute_relaxation(self, x_limit=500):
    method retrieve_from_precompute (line 899) | def retrieve_from_precompute(self, x, flip=False):
    method bound_relax_impl (line 914) | def bound_relax_impl(self, x):
  class BoundSigmoid (line 1014) | class BoundSigmoid(BoundTanh):
    method __init__ (line 1015) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method build_gradient_node (line 1019) | def build_gradient_node(self, grad_upstream):
  class SigmoidGradOp (line 1026) | class SigmoidGradOp(Function):
    method symbolic (line 1028) | def symbolic(_, preact):
    method forward (line 1032) | def forward(ctx, preact):
  class SigmoidGrad (line 1037) | class SigmoidGrad(Module):
    method forward (line 1038) | def forward(self, g, preact):
  class BoundSigmoidGrad (line 1042) | class BoundSigmoidGrad(BoundTanhGrad):
    method __init__ (line 1043) | def __init__(self, attr=None, inputs=None, output_index=0, options=None,
  class BoundAtan (line 1051) | class BoundAtan(BoundTanh):
    method __init__ (line 1052) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method build_gradient_node (line 1057) | def build_gradient_node(self, grad_upstream):
  class AtanGrad (line 1064) | class AtanGrad(Module):
    method forward (line 1065) | def forward(self, g, preact):
  class BoundTan (line 1070) | class BoundTan(BoundAtan):
    method forward (line 1077) | def forward(self, x):
    method _check_bounds (line 1080) | def _check_bounds(self, lower, upper):
    method _init_masks (line 1095) | def _init_masks(self, x):
    method interval_propagate (line 1103) | def interval_propagate(self, *v):
    method bound_relax (line 1110) | def bound_relax(self, x, init=False, dim_opt=None):

FILE: auto_LiRPA/operators/shape.py
  class BoundShape (line 21) | class BoundShape(Bound):
    method __init__ (line 22) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method shape (line 27) | def shape(x):
    method forward (line 30) | def forward(self, x):
    method bound_forward (line 34) | def bound_forward(self, dim_in, x):
    method build_solver (line 37) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...

FILE: auto_LiRPA/operators/slice_concat.py
  class BoundConcat (line 25) | class BoundConcat(Bound):
    method __init__ (line 26) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 32) | def forward(self, *x):  # x is a list of tensors
    method interval_propagate (line 38) | def interval_propagate(self, *v):
    method bound_backward (line 72) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method bound_forward (line 115) | def bound_forward(self, dim_in, *x):
    method build_solver (line 131) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method build_gradient_node (line 134) | def build_gradient_node(self, grad_upstream):
  class BoundSlice (line 146) | class BoundSlice(Bound):
    method __init__ (line 147) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method __repr__ (line 155) | def __repr__(self):
    method _fixup_params (line 166) | def _fixup_params(self, shape, start, end, axes, steps):
    method forward (line 180) | def forward(self, x, start=None, end=None, axes=None, steps=1):
    method interval_propagate (line 192) | def interval_propagate(self, *v):
    method build_solver (line 197) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
    method bound_backward (line 200) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method bound_forward (line 241) | def bound_forward(self, dim_in, *inputs):
    method build_gradient_node (line 263) | def build_gradient_node(self, grad_upstream):
  class BoundSplit (line 275) | class BoundSplit(Bound):
    method __init__ (line 276) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 285) | def forward(self, *x):
    method bound_backward (line 292) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method bound_forward (line 315) | def bound_forward(self, dim_in, *x):
    method build_solver (line 325) | def build_solver(self, *v, model, C=None, model_type="mip", solver_pkg...
  function slice_grad (line 329) | def slice_grad(x, input_shape, start, end, axes, steps):
  class SliceGradOp (line 350) | class SliceGradOp(Function):
    method symbolic (line 356) | def symbolic(_, grad_last, input, start=None, end=None, axes=None, ste...
    method forward (line 363) | def forward(ctx, grad_last, input, start, end, axes, steps):
  class SliceGrad (line 367) | class SliceGrad(Module):
    method __init__ (line 368) | def __init__(self, start, end, axes, steps):
    method forward (line 375) | def forward(self, grad_last, input):
  class BoundSliceGrad (line 381) | class BoundSliceGrad(Bound):
    method __init__ (line 382) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 390) | def forward(self, grad_last, input):
    method bound_backward (line 394) | def bound_backward(self, last_lA, last_uA, *args, **kwargs):
  function concat_grad (line 409) | def concat_grad(x, axis, input_index, *inputs):
  class ConcatGradOp (line 419) | class ConcatGradOp(Function):
    method symbolic (line 421) | def symbolic(_, grad_last, axis, input_index, *inputs):
    method forward (line 426) | def forward(ctx, grad_last, axis, input_index, *inputs):
  class ConcatGrad (line 430) | class ConcatGrad(Module):
    method __init__ (line 431) | def __init__(self, axis, input_index):
    method forward (line 436) | def forward(self, grad_last, *input):
  class BoundConcatGrad (line 440) | class BoundConcatGrad(Bound):
    method __init__ (line 441) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 447) | def forward(self, grad_last, *inputs):
    method bound_backward (line 450) | def bound_backward(self, last_lA, last_uA, *args, **kwargs):

FILE: auto_LiRPA/operators/softmax.py
  class BoundSoftmaxImpl (line 20) | class BoundSoftmaxImpl(nn.Module):
    method __init__ (line 21) | def __init__(self, axis):
    method forward (line 26) | def forward(self, x):
  class BoundSoftmax (line 33) | class BoundSoftmax(Bound):
    method __init__ (line 34) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 43) | def forward(self, x):
    method interval_propagate (line 53) | def interval_propagate(self, *v):

FILE: auto_LiRPA/operators/solver_utils.py
  class DummyGurobipyClass (line 17) | class DummyGurobipyClass:
    method __getattr__ (line 19) | def __getattr__(self, attr):

FILE: auto_LiRPA/operators/tile.py
  class BoundTile (line 21) | class BoundTile(Bound):
    method __init__ (line 22) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 26) | def forward(self, x, repeats):
    method bound_backward (line 29) | def bound_backward(self, last_lA, last_uA, *x, **kwargs):
    method bound_forward (line 51) | def bound_forward(self, dim_in, *x):

FILE: auto_LiRPA/operators/trigonometric.py
  class BoundSin (line 26) | class BoundSin(BoundSShaped):
    method __init__ (line 35) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method d2_act_func (line 55) | def d2_act_func(self, x):
    method _init_opt_parameters_impl (line 58) | def _init_opt_parameters_impl(self, size_spec, name_start):
    method opt_init (line 70) | def opt_init(self):
    method branch_input_domain (line 77) | def branch_input_domain(self, lb, ub):
    method generate_d_lower_upper (line 122) | def generate_d_lower_upper(self, lower, upper):
    method arcsin (line 170) | def arcsin(c):
    method get_intersection (line 188) | def get_intersection(start, end, c, theta=0.):
    method n_crossing (line 196) | def n_crossing(start, end, s):
    method check_bound (line 206) | def check_bound(tangent_point, x):
    method get_lower_left_bound (line 218) | def get_lower_left_bound(xl, steps=20):
    method get_upper_left_bound (line 245) | def get_upper_left_bound(xl, steps=20):
    method get_lower_right_bound (line 272) | def get_lower_right_bound(xu, steps=20):
    method get_upper_right_bound (line 282) | def get_upper_right_bound(xu, steps=20):
    method get_bound_tb (line 290) | def get_bound_tb(self, lb, ub):
    method forward (line 385) | def forward(self, x):
    method interval_propagate (line 388) | def interval_propagate(self, *v):
    method bound_relax_branch (line 405) | def bound_relax_branch(self, lb, ub):
  class BoundCos (line 445) | class BoundCos(BoundSin):
    method __init__ (line 446) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 451) | def forward(self, x):
    method bound_relax (line 454) | def bound_relax(self, x, init=False, dim_opt=None):
  class BoundSec (line 465) | class BoundSec(BoundActivation):
    method __init__ (line 466) | def __init__(self, attr=None, inputs=None, output_index=0, options=None):
    method forward (line 470) | def forward(self, x):
    method bound_relax (line 473) | def bound_relax(self, x, init=False):
    method interval_propagate (line 496) | def interval_propagate(self, *v):
  class SinGradOp (line 507) | class SinGradOp(Function):
    method symbolic (line 509) | def symbolic(_, x):
    method forward (line 513) | def forward(ctx, input):
  class CosGradOp (line 517) | class CosGradOp(Function):
    method symbolic (line 519) | def symbolic(_, x):
    method forward (line 523) | def forward(ctx, input):
  class TanhGradOp (line 527) | class TanhGradOp(Function):
    method symbolic (line 529) | def symbolic(_, x):
    method forward (line 533) | def forward(ctx, input):

FILE: auto_LiRPA/opt_pruner.py
  class OptPruner (line 24) | class OptPruner:
    method __init__ (line 26) | def __init__(self, x, threshold, multi_spec_keep_func, loss_reduction_...
    method prune (line 47) | def prune(self, x, C, ret_l, ret_u, ret, full_l, full_ret_l, full_ret_u,
    method prune_idx (line 110) | def prune_idx(self, idx_mask, idx, x):
    method next_iter (line 126) | def next_iter(self):
    method update_best (line 136) | def update_best(self, full_ret_l, full_ret_u, best_ret):
    method update_ratio (line 151) | def update_ratio(self, full_l, full_ret_l):
    method _prune_x (line 181) | def _prune_x(self, x):
    method _prune_dict_of_lists (line 197) | def _prune_dict_of_lists(self, dict_of_lists, pre_prune_size):
    method _prune_bounds_by_mask (line 210) | def _prune_bounds_by_mask(self, ret_l, ret_u, ret, interm_bounds,
    method get_preserve_mask (line 235) | def get_preserve_mask(self, ret_l):
    method _recover_bounds_to_full_batch (line 257) | def _recover_bounds_to_full_batch(self, ret):
    method cache_full_sized_alpha (line 287) | def cache_full_sized_alpha(self, optimizable_activations: list):
    method recover_full_sized_alpha (line 307) | def recover_full_sized_alpha(self, optimizable_activations: list):
    method clean_full_sized_alpha_cache (line 318) | def clean_full_sized_alpha_cache(self):

FILE: auto_LiRPA/optimize_graph.py
  function _optimize_graph (line 31) | def _optimize_graph(self: 'BoundedModule'):
  function _copy_node_properties (line 51) | def _copy_node_properties(new, ref):
  function merge_sec (line 59) | def merge_sec(model: 'BoundedModule'):
  function div_to_mul (line 76) | def div_to_mul(model: 'BoundedModule'):
  function convert_sqr (line 97) | def convert_sqr(model: 'BoundedModule'):
  function merge_identical_act (line 121) | def merge_identical_act(model: 'BoundedModule'):
  function minmax_to_relu (line 139) | def minmax_to_relu(model: 'BoundedModule'):
  function _pair_row (line 184) | def _pair_row(Ws, bs, Wm, j, atol=1e-8):
  function optimize_relu_relation (line 216) | def optimize_relu_relation(model: 'BoundedModule'):

FILE: auto_LiRPA/optimized_bounds.py
  function opt_reuse (line 98) | def opt_reuse(self: 'BoundedModule'):
  function opt_no_reuse (line 103) | def opt_no_reuse(self: 'BoundedModule'):
  function _set_alpha (line 108) | def _set_alpha(optimizable_activations, parameters, alphas, lr):
  function _set_gammas (line 130) | def _set_gammas(nodes, parameters):
  function _save_ret_first_time (line 144) | def _save_ret_first_time(bounds, best_ret):
  function _to_float64 (line 152) | def _to_float64(self: 'BoundedModule', C, x, aux_reference_bounds, inter...
  function _to_default_dtype (line 169) | def _to_default_dtype(self: 'BoundedModule', x, total_loss, full_ret, ret,
  function _get_idx_mask (line 195) | def _get_idx_mask(idx: int, full_ret_bound: Tensor, best_ret_bound: Tens...
  function _update_best_ret (line 220) | def _update_best_ret(
  function _update_A_dict (line 307) | def _update_A_dict(best_A, full_A, improved_idx, c_mask: Optional[Tensor...
  function _update_optimizable_activations (line 336) | def _update_optimizable_activations(
  function update_best_beta (line 370) | def update_best_beta(self: 'BoundedModule', enable_opt_interm_bounds, be...
  function _get_optimized_bounds (line 396) | def _get_optimized_bounds(
  function init_alpha (line 970) | def init_alpha(self: 'BoundedModule', x, share_alphas=False, method='bac...

FILE: auto_LiRPA/output_constraints.py
  function invprop_enabled (line 26) | def invprop_enabled(self: 'BoundedModule'):
  function invprop_init_infeasible_bounds (line 30) | def invprop_init_infeasible_bounds(self: 'BoundedModule', bound_node, C):
  function invprop_check_infeasible_bounds (line 43) | def invprop_check_infeasible_bounds(self: 'BoundedModule', lb, ub):
  function backward_general_invprop (line 56) | def backward_general_invprop(
  function backward_general_with_output_constraint (line 140) | def backward_general_with_output_constraint(

FILE: auto_LiRPA/parse_graph.py
  function get_node_name (line 32) | def get_node_name(node):
  function get_node_attribute (line 35) | def get_node_attribute(node, attribute_name):
  function parse_graph (line 43) | def parse_graph(graph, inputs, params):
  function _get_jit_params (line 140) | def _get_jit_params(module, param_exclude, param_include):
  function get_output_template (line 161) | def get_output_template(out):
  function parse_source (line 178) | def parse_source(node):
  function update_debug_names (line 191) | def update_debug_names(trace_graph):
  function parse_module (line 203) | def parse_module(module, inputs, param_exclude=".*AuxLogits.*", param_in...

FILE: auto_LiRPA/patches.py
  function insert_zeros (line 22) | def insert_zeros(image, s):
  function remove_zeros (line 58) | def remove_zeros(image, s, remove_zero_start_idx=(0,0)):
  function unify_shape (line 78) | def unify_shape(shape):
  function simplify_shape (line 94) | def simplify_shape(shape):
  function is_shape_used (line 111) | def is_shape_used(shape, expected=0):
  class Patches (line 118) | class Patches:
    method __init__ (line 127) | def __init__(
    method __add__ (line 147) | def __add__(self, other):
    method __str__ (line 190) | def __str__(self):
    method device (line 199) | def device(self):
    method create_similar (line 209) | def create_similar(self, patches=None, stride=None, padding=None, iden...
    method clone (line 232) | def clone(self):
    method detach (line 235) | def detach(self):
    method to_matrix (line 254) | def to_matrix(self, input_shape):
    method simplify (line 261) | def simplify(self):
    method matmul (line 285) | def matmul(self, input, patch_abs=False, input_shape=None):
    method create_padding (line 322) | def create_padding(self, output_shape):
  function compute_patches_stride_padding (line 338) | def compute_patches_stride_padding(input_shape, patches_padding, patches...
  function patches_to_matrix (line 373) | def patches_to_matrix(pieces, input_shape, stride, padding, output_shape...
  function check_patch_biases (line 448) | def check_patch_biases(lb, ub, lower_b, upper_b):
  function inplace_unfold (line 465) | def inplace_unfold(image, kernel_size, stride=1, padding=0, inserted_zer...
  function maybe_unfold_patches (line 519) | def maybe_unfold_patches(d_tensor, last_A, alpha_lookup_idx=None):
  function create_valid_mask (line 634) | def create_valid_mask(output_shape, device, dtype, kernel_size, stride, ...

FILE: auto_LiRPA/perturbations.py
  class Perturbation (line 28) | class Perturbation:
    method __init__ (line 42) | def __init__(self):
    method set_eps (line 45) | def set_eps(self, eps):
    method concretize (line 48) | def concretize(self, x, A, sign=-1, aux=None):
    method init (line 66) | def init(self, x, aux=None, forward=False):
  class PerturbationL0Norm (line 88) | class PerturbationL0Norm(Perturbation):
    method __init__ (line 94) | def __init__(self, eps, x_L=None, x_U=None, ratio=1.0):
    method concretize (line 100) | def concretize(self, x, A, sign=-1, aux=None):
    method init (line 130) | def init(self, x, aux=None, forward=False):
    method __repr__ (line 144) | def __repr__(self):
  class PerturbationLpNorm (line 148) | class PerturbationLpNorm(Perturbation):
    method __init__ (line 150) | def __init__(self, eps=0, norm=np.inf, x_L=None, x_U=None, eps_min=0,
    method get_input_bounds (line 231) | def get_input_bounds(self, x, A):
    method get_constraints (line 246) | def get_constraints(self, A):
    method concretize_matrix (line 254) | def concretize_matrix(self, x, A, sign, constraints=None):
    method concretize_patches (line 294) | def concretize_patches(self, x, A, sign):
    method concretize (line 340) | def concretize(self, x, A, sign=-1, constraints=None, aux=None):
    method init_sparse_linf (line 354) | def init_sparse_linf(self, x, x_L, x_U):
    method init (line 408) | def init(self, x, aux=None, forward=False):
    method add_infeasible_batches (line 447) | def add_infeasible_batches(self, infeasible_batches):
    method add_objective_indices (line 480) | def add_objective_indices(self, objective_indices):
    method constraints_enable (line 485) | def constraints_enable(self):
    method constraints_enable (line 492) | def constraints_enable(self, enable: bool):
    method constraints_enable (line 496) | def constraints_enable(self):
    method use_grad (line 500) | def use_grad(self):
    method use_grad (line 507) | def use_grad(self, use_grad: bool):
    method use_grad (line 511) | def use_grad(self):
    method add_aux_bounds (line 514) | def add_aux_bounds(self, aux_lb, aux_ub):
    method clear_aux_bounds (line 518) | def clear_aux_bounds(self):
    method reset_constraints (line 522) | def reset_constraints(self, constraints, decision_thresh):
    method __repr__ (line 535) | def __repr__(self):
  class PerturbationLinear (line 545) | class PerturbationLinear(Perturbation):
    method __init__ (line 560) | def __init__(self, lower_A, upper_A, lower_b, upper_b, input_lb, input...
    method concretize (line 577) | def concretize(self, x, A, sign=-1, aux=None):
    method init (line 597) | def init(self, x, aux=None, forward=False):
  class PerturbationSynonym (line 604) | class PerturbationSynonym(Perturbation):
    method __init__ (line 605) | def __init__(self, budget, eps=1.0, use_simple=False):
    method __repr__ (line 614) | def __repr__(self):
    method _load_synonyms (line 618) | def _load_synonyms(self, path='data/synonyms.json'):
    method set_train (line 623) | def set_train(self, train):
    method concretize (line 626) | def concretize(self, x, A, sign, aux):
    method init (line 710) | def init(self, x, aux=None, forward=False):
    method _build_substitution (line 794) | def _build_substitution(self, batch):

FILE: auto_LiRPA/solver_module.py
  function build_solver_module (line 24) | def build_solver_module(self: 'BoundedModule', x=None, C=None, interm_bo...
  function _build_solver_general (line 81) | def _build_solver_general(self: 'BoundedModule', node: Bound, C=None, mo...
  function _reset_solver_vars (line 104) | def _reset_solver_vars(self: 'BoundedModule', node: Bound, iteration=True):
  function _reset_solver_model (line 112) | def _reset_solver_model(self: 'BoundedModule'):
  function _build_solver_input (line 117) | def _build_solver_input(self: 'BoundedModule', node):

FILE: auto_LiRPA/tools.py
  function visualize (line 27) | def visualize(self: 'BoundedModule', output_path, print_bounds=False):
  function sanitize_graphviz_name (line 96) | def sanitize_graphviz_name(name):

FILE: auto_LiRPA/utils.py
  function onehotc_to_dense (line 48) | def onehotc_to_dense(one_hot_c: OneHotC, dtype: torch.dtype) -> torch.Te...
  function reduction_str2func (line 72) | def reduction_str2func(reduction_func):
  function stop_criterion_placeholder (line 87) | def stop_criterion_placeholder(threshold=0):
  function stop_criterion_min (line 90) | def stop_criterion_min(threshold=0):
  function stop_criterion_all (line 93) | def stop_criterion_all(threshold=0):
  function stop_criterion_max (line 99) | def stop_criterion_max(threshold=0):
  function stop_criterion_batch (line 102) | def stop_criterion_batch(threshold=0):
  function stop_criterion_batch_any (line 107) | def stop_criterion_batch_any(threshold=0):
  function stop_criterion_general (line 115) | def stop_criterion_general(or_spec_size, threshold=0):
  function stop_criterion_batch_topk (line 139) | def stop_criterion_batch_topk(threshold=0, k=1314):
  function multi_spec_keep_func_all (line 143) | def multi_spec_keep_func_all(x):
  class MultiAverageMeter (line 155) | class MultiAverageMeter(object):
    method __init__ (line 157) | def __init__(self):
    method reset (line 160) | def reset(self):
    method set_batch_size (line 166) | def set_batch_size(self, batch_size):
    method update (line 169) | def update(self, key, val, n=None):
    method last (line 180) | def last(self, key):
    method avg (line 183) | def avg(self, key):
    method __repr__ (line 189) | def __repr__(self):
  class MultiTimer (line 196) | class MultiTimer(object):
    method __init__ (line 198) | def __init__(self):
    method reset (line 200) | def reset(self):
    method start (line 203) | def start(self, key):
    method stop (line 207) | def stop(self, key):
    method total (line 212) | def total(self, key):
    method __repr__ (line 214) | def __repr__(self):
  class Flatten (line 221) | class Flatten(nn.Flatten):
  class Unflatten (line 229) | class Unflatten(nn.Module):
    method __init__ (line 230) | def __init__(self, wh):
    method forward (line 233) | def forward(self, x):
  class Max (line 237) | class Max(nn.Module):
    method __init__ (line 239) | def __init__(self):
    method forward (line 242) | def forward(self, x, y):
  class Min (line 246) | class Min(nn.Module):
    method __init__ (line 248) | def __init__(self):
    method forward (line 251) | def forward(self, x, y):
  function scale_gradients (line 255) | def scale_gradients(optimizer, gradient_accumulation_steps, grad_clip=No...
  function unpack_inputs (line 268) | def unpack_inputs(inputs, device=None):
  function isnan (line 282) | def isnan(x):
  function prod (line 288) | def prod(x):
  function batched_index_select (line 292) | def batched_index_select(input, dim, index):
  function get_spec_matrix (line 313) | def get_spec_matrix(X, y, num_classes):
  function unravel_index (line 322) | def unravel_index(
  class AutoBatchSize (line 348) | class AutoBatchSize:
    method __init__ (line 349) | def __init__(self, init_batch_size, device, vram_ratio=0.9, enable=True):
    method record_actual_batch_size (line 356) | def record_actual_batch_size(self, actual_batch_size):
    method update (line 363) | def update(self):
  function sync_params (line 384) | def sync_params(model_ori: torch.nn.Module,
  function reduce_broadcast_dims (line 408) | def reduce_broadcast_dims(A, target_shape, left_extra_dims=1):
  function matmul_maybe_batched (line 453) | def matmul_maybe_batched(a: torch.Tensor, b: torch.Tensor, both_batched:...
  function transfer (line 460) | def transfer(tensor, device=None, dtype=None, non_blocking=False):
  function clone_sub_A_dict (line 470) | def clone_sub_A_dict(A_dict, out_in_keys: Tuple):
  function clone_full_A_dict (line 504) | def clone_full_A_dict(A_dict):

FILE: auto_LiRPA/wrapper.py
  class CrossEntropyWrapper (line 20) | class CrossEntropyWrapper(nn.Module):
    method __init__ (line 21) | def __init__(self, model):
    method forward (line 25) | def forward(self, x, labels):
  class CrossEntropyWrapperMultiInput (line 30) | class CrossEntropyWrapperMultiInput(nn.Module):
    method __init__ (line 31) | def __init__(self, model):
    method forward (line 35) | def forward(self, labels, *x):

FILE: doc/conf.py
  function linkcode_resolve (line 65) | def linkcode_resolve(domain, info):

FILE: examples/language/Transformer/Transformer.py
  class Transformer (line 30) | class Transformer(nn.Module):
    method __init__ (line 31) | def __init__(self, args, data_train):
    method save (line 69) | def save(self, epoch):
    method build_optimizer (line 79) | def build_optimizer(self):
    method train (line 87) | def train(self):
    method eval (line 91) | def eval(self):
    method get_input (line 95) | def get_input(self, batch):
    method forward (line 110) | def forward(self, batch):

FILE: examples/language/Transformer/modeling.py
  class BertLayerNorm (line 25) | class BertLayerNorm(nn.Module):
    method __init__ (line 26) | def __init__(self, hidden_size, eps=1e-12):
    method forward (line 32) | def forward(self, x):
  class BertLayerNormNoVar (line 38) | class BertLayerNormNoVar(nn.Module):
    method __init__ (line 39) | def __init__(self, hidden_size, eps=1e-12):
    method forward (line 45) | def forward(self, x):
  class BertEmbeddings (line 50) | class BertEmbeddings(nn.Module):
    method __init__ (line 53) | def __init__(self, config, glove=None, vocab=None):
    method forward (line 61) | def forward(self, input_ids, token_type_ids=None):
  class BertSelfOutput (line 78) | class BertSelfOutput(nn.Module):
    method __init__ (line 79) | def __init__(self, config):
    method forward (line 89) | def forward(self, hidden_states, input_tensor):
  class BertAttention (line 100) | class BertAttention(nn.Module):
    method __init__ (line 101) | def __init__(self, config, input_size):
    method forward (line 106) | def forward(self, input_tensor, attention_mask):
  class BertOutput (line 112) | class BertOutput(nn.Module):
    method __init__ (line 113) | def __init__(self, config):
    method forward (line 123) | def forward(self, hidden_states, input_tensor):
  class BertLayer (line 133) | class BertLayer(nn.Module):
    method __init__ (line 134) | def __init__(self, config, layer_id):
    method forward (line 141) | def forward(self, hidden_states, attention_mask):
  class BertEncoder (line 148) | class BertEncoder(nn.Module):
    method __init__ (line 149) | def __init__(self, config):
    method forward (line 153) | def forward(self, hidden_states, attention_mask, output_all_encoded_la...
  class BertPooler (line 163) | class BertPooler(nn.Module):
    method __init__ (line 164) | def __init__(self, config):
    method forward (line 169) | def forward(self, hidden_states):
  class BertModelFromEmbeddings (line 177) | class BertModelFromEmbeddings(BertPreTrainedModel):
    method __init__ (line 178) | def __init__(self, config):
    method forward (line 184) | def forward(self, embeddings, extended_attention_mask):
  class BertForSequenceClassificationFromEmbeddings (line 190) | class BertForSequenceClassificationFromEmbeddings(BertPreTrainedModel):
    method __init__ (line 191) | def __init__(self, config, num_labels=2):
    method forward (line 207) | def forward(self, embeddings, extended_attention_mask):
  class BertForSequenceClassification (line 221) | class BertForSequenceClassification(BertPreTrainedModel):
    method __init__ (line 222) | def __init__(self, config, num_labels=2, glove=None, vocab=None):
    method forward (line 231) | def forward(self, input_ids, token_type_ids=None, attention_mask=None,...

FILE: examples/language/Transformer/utils.py
  class InputExample (line 19) | class InputExample(object):
    method __init__ (line 20) | def __init__(self, guid, text_a, text_b=None, label=None):
  class InputFeatures (line 26) | class InputFeatures(object):
    method __init__ (line 27) | def __init__(self, input_ids, input_mask, segment_ids, label_id, tokens):
  function convert_examples_to_features (line 34) | def convert_examples_to_features(examples, label_list, max_seq_length,

FILE: examples/language/data_utils.py
  function load_data_sst (line 5) | def load_data_sst():
  function load_data (line 12) | def load_data(dataset):
  function clean_data (line 18) | def clean_data(data):
  function get_batches (line 21) | def get_batches(data, batch_size):

FILE: examples/language/language_utils.py
  function build_vocab (line 4) | def build_vocab(data_train, min_word_freq, dump=False, include=[]):
  function tokenize (line 31) | def tokenize(batch, vocab, max_seq_length, drop_unk=False):
  function token_to_id (line 47) | def token_to_id(tokens, vocab):

FILE: examples/language/lstm.py
  class LSTMFromEmbeddings (line 9) | class LSTMFromEmbeddings(nn.Module):
    method __init__ (line 10) | def __init__(self, args, vocab_size):
    method forward (line 27) | def forward(self, embeddings, mask):
  class LSTM (line 46) | class LSTM(nn.Module):
    method __init__ (line 47) | def __init__(self, args, data_train):
    method save (line 79) | def save(self, epoch):
    method build_optimizer (line 88) | def build_optimizer(self):
    method get_input (line 97) | def get_input(self, batch):
    method train (line 122) | def train(self):
    method eval (line 125) | def eval(self):

FILE: examples/language/oracle.py
  function oracle (line 6) | def oracle(args, model, ptb, data, type):

FILE: examples/language/preprocess/pre_compute_lm_scores.py
  function parse_args (line 18) | def parse_args():
  function main (line 26) | def main():

FILE: examples/language/preprocess/preprocess_sst.py
  function load_data_sst (line 3) | def load_data_sst():
  function read_scores (line 48) | def read_scores(split):

FILE: examples/language/train.py
  function step (line 143) | def step(model, ptb, batch, eps=1.0, train=False):
  function train (line 233) | def train(epoch, batches, type):
  function main (line 292) | def main():

FILE: examples/sequence/data_utils.py
  function load_data (line 5) | def load_data():
  function get_batches (line 13) | def get_batches(data, batch_size):

FILE: examples/sequence/lstm.py
  class LSTMCore (line 7) | class LSTMCore(nn.Module):
    method __init__ (line 8) | def __init__(self, args):
    method forward (line 19) | def forward(self, X):
  class LSTM (line 31) | class LSTM(nn.Module):
    method __init__ (line 32) | def __init__(self, args):
    method save (line 52) | def save(self, epoch):
    method build_optimizer (line 63) | def build_optimizer(self):
    method get_input (line 70) | def get_input(self, batch):
    method train (line 75) | def train(self):
    method eval (line 78) | def eval(self):

FILE: examples/sequence/train.py
  function step (line 32) | def step(model, ptb, batch, eps=args.eps, train=False):
  function train (line 64) | def train(epoch):
  function test (line 92) | def test(epoch, batches):

FILE: examples/simple/invprop.py
  class simple_model (line 11) | class simple_model(torch.nn.Module):
    method __init__ (line 15) | def __init__(self):
    method forward (line 21) | def forward(self, x):

FILE: examples/simple/lp_full.py
  function build_C (line 33) | def build_C(label, classes):

FILE: examples/simple/mip_lp_solver.py
  class mnist_model (line 22) | class mnist_model(nn.Module):
    method __init__ (line 23) | def __init__(
    method forward (line 32) | def forward(self, x):

FILE: examples/simple/toy.py
  class simple_model (line 9) | class simple_model(torch.nn.Module):
    method __init__ (line 13) | def __init__(self):
    method forward (line 19) | def forward(self, x):

FILE: examples/vision/bound_option.py
  function mnist_model (line 16) | def mnist_model():

FILE: examples/vision/cifar_training.py
  function get_exp_module (line 20) | def get_exp_module(bounded_module):
  function Train (line 60) | def Train(model, t, loader, eps_scheduler, norm, train, opt, bound_type,...
  function main (line 190) | def main(args):

FILE: examples/vision/custom_op.py
  class PlusConstantOp (line 16) | class PlusConstantOp(torch.autograd.Function):
    method symbolic (line 18) | def symbolic(g, x, const):
    method forward (line 28) | def forward(ctx, x, const):
  class PlusConstant (line 35) | class PlusConstant(nn.Module):
    method __init__ (line 36) | def __init__(self, const=1):
    method forward (line 40) | def forward(self, x):
  class BoundPlusConstant (line 45) | class BoundPlusConstant(Bound):
    method __init__ (line 46) | def __init__(self, attr, inputs, output_index, options):
    method forward (line 51) | def forward(self, x):
    method bound_backward (line 54) | def bound_backward(self, last_lA, last_uA, x, *args, **kwargs):
    method interval_propagate (line 73) | def interval_propagate(self, *v):

FILE: examples/vision/data/ImageNet64/imagenet_data_loader.py
  class DatasetDownsampledImageNet (line 7) | class DatasetDownsampledImageNet():
    method __init__ (line 8) | def __init__(self):
    method load_data (line 19) | def load_data(self, data_path, img_size=64, count=0., fname=''):

FILE: examples/vision/datasets.py
  function get_stats (line 9) | def get_stats(loader):
  function mnist_loaders (line 28) | def mnist_loaders(dataset, batch_size, shuffle_train = True, shuffle_tes...
  function cifar_loaders (line 61) | def cifar_loaders(batch_size, shuffle_train = True, shuffle_test = False...
  function svhn_loaders (line 107) | def svhn_loaders(batch_size, shuffle_train = True, shuffle_test = False,...
  function load_data (line 154) | def load_data(data, batch_size):

FILE: examples/vision/imagenet_training.py
  function get_exp_module (line 17) | def get_exp_module(bounded_module):
  function Train (line 59) | def Train(model, t, loader, eps_scheduler, norm, train, opt, bound_type,...
  function main (line 193) | def main(args):

FILE: examples/vision/jacobian.py
  function build_model (line 18) | def build_model(in_ch=3, in_dim=32):
  function example_jacobian (line 30) | def example_jacobian(model_ori, x0, bound_opts, device):
  function example_local_lipschitz (line 68) | def example_local_lipschitz(model_ori, x0, bound_opts, device):
  function example_jvp (line 110) | def example_jvp(model_ori, x0, bound_opts, device):
  function compute_jacobians (line 146) | def compute_jacobians(model_ori, x0, bound_opts=None, device='cpu'):

FILE: examples/vision/models/densenet.py
  class Bottleneck (line 11) | class Bottleneck(nn.Module):
    method __init__ (line 12) | def __init__(self, in_planes, growth_rate):
    method forward (line 19) | def forward(self, x):
  class Transition (line 28) | class Transition(nn.Module):
    method __init__ (line 29) | def __init__(self, in_planes, out_planes):
    method forward (line 34) | def forward(self, x):
  class DenseNet (line 40) | class DenseNet(nn.Module):
    method __init__ (line 41) | def __init__(self, block, nblocks, growth_rate=12, reduction=0.5, num_...
    method _make_dense_layers (line 74) | def _make_dense_layers(self, block, in_planes, nblock):
    method forward (line 81) | def forward(self, x):
  function Densenet_cifar_32 (line 93) | def Densenet_cifar_32(in_ch=3, in_dim=32):

FILE: examples/vision/models/densenet_imagenet.py
  class Bottleneck (line 11) | class Bottleneck(nn.Module):
    method __init__ (line 12) | def __init__(self, in_planes, growth_rate):
    method forward (line 19) | def forward(self, x):
  class Transition (line 28) | class Transition(nn.Module):
    method __init__ (line 29) | def __init__(self, in_planes, out_planes):
    method forward (line 34) | def forward(self, x):
  class DenseNet (line 40) | class DenseNet(nn.Module):
    method __init__ (line 41) | def __init__(self, block, nblocks, growth_rate=12, reduction=0.5, num_...
    method _make_dense_layers (line 74) | def _make_dense_layers(self, block, in_planes, nblock):
    method forward (line 81) | def forward(self, x):
  function Densenet_imagenet (line 94) | def Densenet_imagenet(in_ch=3, in_dim=56):

FILE: examples/vision/models/densenet_no_bn.py
  class Bottleneck (line 12) | class Bottleneck(nn.Module):
    method __init__ (line 13) | def __init__(self, in_planes, growth_rate):
    method forward (line 20) | def forward(self, x):
  class Transition (line 29) | class Transition(nn.Module):
    method __init__ (line 30) | def __init__(self, in_planes, out_planes):
    method forward (line 35) | def forward(self, x):
  class DenseNet (line 41) | class DenseNet(nn.Module):
    method __init__ (line 42) | def __init__(self, block, nblocks, growth_rate=12, reduction=0.5, num_...
    method _make_dense_layers (line 75) | def _make_dense_layers(self, block, in_planes, nblock):
    method forward (line 82) | def forward(self, x):
  function Densenet_cifar_wobn (line 94) | def Densenet_cifar_wobn(in_ch=3, in_dim=56):

FILE: examples/vision/models/feedforward.py
  class cnn_4layer (line 11) | class cnn_4layer(nn.Module):
    method __init__ (line 12) | def __init__(self, in_ch, in_dim, width=2, linear_size=256):
    method forward (line 19) | def forward(self, x):
  class mlp_2layer (line 29) | class mlp_2layer(nn.Module):
    method __init__ (line 30) | def __init__(self, in_ch, in_dim, width=1):
    method forward (line 35) | def forward(self, x):
  class mlp_3layer (line 42) | class mlp_3layer(nn.Module):
    method __init__ (line 43) | def __init__(self, in_ch, in_dim, width=1):
    method forward (line 49) | def forward(self, x):
  class mlp_3layer_weight_perturb (line 57) | class mlp_3layer_weight_perturb(nn.Module):
    method __init__ (line 58) | def __init__(self, in_ch=1, in_dim=28, width=1, pert_weight=True, pert...
    method forward (line 77) | def forward(self, x):
  class mlp_5layer (line 85) | class mlp_5layer(nn.Module):
    method __init__ (line 86) | def __init__(self, in_ch, in_dim, width=1):
    method forward (line 94) | def forward(self, x):
  function cnn_7layer (line 105) | def cnn_7layer(in_ch=3, in_dim=32, width=64, linear_size=512):
  function cnn_7layer_bn (line 124) | def cnn_7layer_bn(in_ch=3, in_dim=32, width=64, linear_size=512):
  function cnn_7layer_bn_imagenet (line 148) | def cnn_7layer_bn_imagenet(in_ch=3, in_dim=32, width=64, linear_size=512):
  function cnn_6layer (line 172) | def cnn_6layer(in_ch, in_dim, width=32, linear_size=256):

FILE: examples/vision/models/mobilenet.py
  class Block (line 11) | class Block(nn.Module):
    method __init__ (line 13) | def __init__(self, in_planes, out_planes, expansion, stride):
    method forward (line 32) | def forward(self, x):
  class MobileNetV2 (line 40) | class MobileNetV2(nn.Module):
    method __init__ (line 50) | def __init__(self, num_classes=10):
    method _make_layers (line 60) | def _make_layers(self, in_planes):
    method forward (line 69) | def forward(self, x):

FILE: examples/vision/models/resnet.py
  class Dense (line 10) | class Dense(nn.Module):
    method __init__ (line 11) | def __init__(self, *Ws):
    method forward (line 17) | def forward(self, *xs):
  class DenseSequential (line 23) | class DenseSequential(nn.Sequential):
    method forward (line 24) | def forward(self, x):
  function model_resnet (line 34) | def model_resnet(in_ch=3, in_dim=32, width=1, mult=16, N=1):

FILE: examples/vision/models/resnet18.py
  class BasicBlock (line 14) | class BasicBlock(nn.Module):
    method __init__ (line 17) | def __init__(self, in_planes, planes, stride=1):
    method forward (line 34) | def forward(self, x):
  class Bottleneck (line 42) | class Bottleneck(nn.Module):
    method __init__ (line 45) | def __init__(self, in_planes, planes, stride=1):
    method forward (line 64) | def forward(self, x):
  class ResNet (line 73) | class ResNet(nn.Module):
    method __init__ (line 74) | def __init__(self, block, num_blocks, num_classes=10, in_planes=64):
    method _make_layer (line 87) | def _make_layer(self, block, planes, num_blocks, stride):
    method forward (line 95) | def forward(self, x):
  function ResNet18 (line 106) | def ResNet18(in_planes=64):

FILE: examples/vision/models/resnext.py
  class Block (line 10) | class Block(nn.Module):
    method __init__ (line 14) | def __init__(self, in_planes, cardinality=32, bottleneck_width=4, stri...
    method forward (line 31) | def forward(self, x):
  class ResNeXt (line 43) | class ResNeXt(nn.Module):
    method __init__ (line 44) | def __init__(self, num_blocks, cardinality, bottleneck_width, num_clas...
    method _make_layer (line 60) | def _make_layer(self, num_blocks, stride):
    method forward (line 70) | def forward(self, x):
  function ResNeXt29_2x64d (line 81) | def ResNeXt29_2x64d():
  function ResNeXt29_4x64d (line 84) | def ResNeXt29_4x64d():
  function ResNeXt29_8x64d (line 87) | def ResNeXt29_8x64d():
  function ResNeXt29_32x4d (line 90) | def ResNeXt29_32x4d():
  function ResNeXt_cifar (line 93) | def ResNeXt_cifar(in_ch=3, in_dim=32):

FILE: examples/vision/models/resnext_imagenet64.py
  class Block (line 10) | class Block(nn.Module):
    method __init__ (line 14) | def __init__(self, in_planes, cardinality=32, bottleneck_width=4, stri...
    method forward (line 31) | def forward(self, x):
  class ResNeXt (line 43) | class ResNeXt(nn.Module):
    method __init__ (line 44) | def __init__(self, num_blocks, cardinality, bottleneck_width, num_clas...
    method _make_layer (line 60) | def _make_layer(self, num_blocks, stride):
    method forward (line 70) | def forward(self, x):
  function ResNeXt_imagenet64 (line 80) | def ResNeXt_imagenet64():

FILE: examples/vision/models/vnncomp_resnet.py
  class BasicBlock (line 6) | class BasicBlock(nn.Module):
    method __init__ (line 9) | def __init__(self, in_planes, planes, stride=1, bn=True, kernel=3):
    method forward (line 53) | def forward(self, x):
  class ResNet5 (line 66) | class ResNet5(nn.Module):
    method __init__ (line 67) | def __init__(self, block, num_blocks=2, num_classes=10, in_planes=64, ...
    method _make_layer (line 85) | def _make_layer(self, block, planes, num_blocks, stride, bn, kernel):
    method forward (line 93) | def forward(self, x):
  class ResNet9 (line 110) | class ResNet9(nn.Module):
    method __init__ (line 111) | def __init__(self, block, num_blocks=2, num_classes=10, in_planes=64, ...
    method _make_layer (line 130) | def _make_layer(self, block, planes, num_blocks, stride, bn, kernel):
    method forward (line 138) | def forward(self, x):
  function resnet2b (line 156) | def resnet2b():
  function resnet4b (line 159) | def resnet4b():

FILE: examples/vision/models/wide_resnet_cifar.py
  function conv3x3 (line 10) | def conv3x3(in_planes, out_planes, stride=1):
  function conv_init (line 13) | def conv_init(m):
  class wide_basic (line 22) | class wide_basic(nn.Module):
    method __init__ (line 23) | def __init__(self, in_planes, planes, dropout_rate, stride=1, use_bn=F...
    method forward (line 41) | def forward(self, x):
  class Wide_ResNet (line 56) | class Wide_ResNet(nn.Module):
    method __init__ (line 57) | def __init__(self, depth, widen_factor, dropout_rate, num_classes, use...
    method _wide_layer (line 82) | def _wide_layer(self, block, planes, num_blocks, dropout_rate, stride):
    method forward (line 92) | def forward(self, x):
  function wide_resnet_cifar (line 106) | def wide_resnet_cifar(in_ch=3, in_dim=32):
  function wide_resnet_cifar_bn (line 109) | def wide_resnet_cifar_bn(in_ch=3, in_dim=32):
  function wide_resnet_cifar_bn_wo_pooling (line 112) | def wide_resnet_cifar_bn_wo_pooling(in_ch=3, in_dim=32): # 1113M, 21M
  function wide_resnet_cifar_bn_wo_pooling_dropout (line 115) | def wide_resnet_cifar_bn_wo_pooling_dropout(in_ch=3, in_dim=32): # 1113M...

FILE: examples/vision/models/wide_resnet_imagenet64.py
  function conv3x3 (line 9) | def conv3x3(in_planes, out_planes, stride=1):
  function conv_init (line 12) | def conv_init(m):
  class wide_basic (line 21) | class wide_basic(nn.Module):
    method __init__ (line 22) | def __init__(self, in_planes, planes, dropout_rate, stride=1):
    method forward (line 36) | def forward(self, x):
  class Wide_ResNet (line 44) | class Wide_ResNet(nn.Module):
    method __init__ (line 45) | def __init__(self, depth, widen_factor, dropout_rate, num_classes,
    method _wide_layer (line 64) | def _wide_layer(self, block, planes, num_blocks, dropout_rate, stride):
    method forward (line 74) | def forward(self, x):
  function wide_resnet_imagenet64 (line 86) | def wide_resnet_imagenet64(in_ch=3, in_dim=56, in_planes=16, widen_facto...
  function wide_resnet_imagenet64_1000class (line 89) | def wide_resnet_imagenet64_1000class(in_ch=3, in_dim=56, in_planes=16, w...

FILE: examples/vision/save_intermediate_bound.py
  function mnist_model (line 12) | def mnist_model():

FILE: examples/vision/simple_training.py
  function Train (line 52) | def Train(model, t, loader, eps_scheduler, norm, train, opt, bound_type,...
  function main (line 149) | def main(args):

FILE: examples/vision/simple_verification.py
  function mnist_model (line 17) | def mnist_model():
  function concretize_bound (line 97) | def concretize_bound(A, bias, xL, xU, upper: bool):

FILE: examples/vision/tinyimagenet_training.py
  function get_exp_module (line 18) | def get_exp_module(bounded_module):
  function Train (line 61) | def Train(model, t, loader, eps_scheduler, norm, train, opt, bound_type,...
  function main (line 195) | def main(args):

FILE: examples/vision/verify_two_node.py
  class cnn_MNIST (line 16) | class cnn_MNIST(nn.Module):
    method __init__ (line 17) | def __init__(self):
    method forward (line 24) | def forward(self, x, y):

FILE: examples/vision/weight_perturbation_training.py
  function get_exp_module (line 28) | def get_exp_module(bounded_module):
  function Train (line 69) | def Train(model, t, loader, eps_scheduler, norm, train, opt, bound_type,...
  function main (line 189) | def main(args):

FILE: tests/test_1d_activation.py
  class test_model (line 14) | class test_model(nn.Module):
    method __init__ (line 15) | def __init__(self, act_func):
    method forward (line 19) | def forward(self, x):
  function pow_2 (line 22) | def pow_2(x):
  function pow_3 (line 25) | def pow_3(x):
  class GELUOp (line 28) | class GELUOp(torch.autograd.Function):
    method symbolic (line 30) | def symbolic(g, x):
    method forward (line 34) | def forward(ctx, x):
  function GELU (line 37) | def GELU(x):
  function gen_hardtanh (line 40) | def gen_hardtanh(min_val, max_val):
  function tanhgrad (line 45) | def tanhgrad(x):
  function sigmoidgrad (line 48) | def sigmoidgrad(x):
  class Test1DActivation (line 52) | class Test1DActivation(TestCase):
    method __init__ (line 53) | def __init__(self, methodName='runTest', device=DEFAULT_DEVICE, dtype=...
    method create_test (line 56) | def create_test(self, act_func, low, high, ntests=1000, nsamples=1000,
    method test_tan (line 135) | def test_tan(self):
    method test_acts (line 147) | def test_acts(self):

FILE: tests/test_2d_activation.py
  class test_model (line 11) | class test_model(nn.Module):
    method __init__ (line 12) | def __init__(self, act_func):
    method forward (line 16) | def forward(self, x, y):
  function mul (line 20) | def mul(x, y):
  class Test2DActivation (line 24) | class Test2DActivation(TestCase):
    method __init__ (line 25) | def __init__(self, methodName='runTest', device=DEFAULT_DEVICE, dtype=...
    method create_test (line 28) | def create_test(self, act_func, low_x, high_x, low_y, high_y,
    method test_max (line 114) | def test_max(self):
    method test_min (line 118) | def test_min(self):
    method test_mul (line 122) | def test_mul(self):

FILE: tests/test_avgpool.py
  function ff (line 9) | def ff(num_conv=2, num_mlp_only=None, pooling=False, activation="ReLU",
  function synthetic_net (line 55) | def synthetic_net(input_ch, input_dim, **kwargs):
  function synthetic_4c2f_pool (line 59) | def synthetic_4c2f_pool(input_ch, input_dim, **kwargs):
  class TestAvgPool (line 63) | class TestAvgPool(TestCase):
    method __init__ (line 64) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 69) | def test(self):

FILE: tests/test_bound_ops.py
  class Dummy (line 8) | class Dummy:
    method __init__ (line 10) | def __init__(self, lower, upper=None, perturbed=False):
  class TestBoundOp (line 17) | class TestBoundOp(TestCase):
    method __init__ (line 18) | def __init__(self, methodName='runTest', generate=False,
    method test (line 24) | def test(self):

FILE: tests/test_branching_heuristics.py
  function test_branching_heuristics (line 10) | def test_branching_heuristics():

FILE: tests/test_clip_domains.py
  function setup_module (line 29) | def setup_module(module):
  function setup_function (line 43) | def setup_function(function):
  function _tensor (line 51) | def _tensor(x):
  function test_case_one_one (line 54) | def test_case_one_one():
  function test_case_one_two (line 83) | def test_case_one_two():
  function test_case_one_three (line 112) | def test_case_one_three():
  function test_case_one_four (line 142) | def test_case_one_four():
  function test_case_two_one (line 161) | def test_case_two_one():
  function test_case_two_two (line 196) | def test_case_two_two():
  function test_case_two_three (line 230) | def test_case_two_three():
  function concretize_bounds (line 267) | def concretize_bounds(
  function setup_test_matrices (line 313) | def setup_test_matrices(
  function create_batch_copies (line 352) | def create_batch_copies(
  function random_setup_generator (line 381) | def random_setup_generator(

FILE: tests/test_constant.py
  class cnn_MNIST (line 11) | class cnn_MNIST(nn.Module):
    method __init__ (line 12) | def __init__(self):
    method forward (line 19) | def forward(self, x):
  class TestConstant (line 28) | class TestConstant(TestCase):
    method __init__ (line 30) | def __init__(self, methodName='runTest', generate=False,
    method test (line 37) | def test(self):

FILE: tests/test_constrained_concretize.py
  class ConstrainedConcretizeModel (line 9) | class ConstrainedConcretizeModel(torch.nn.Module):
    method __init__ (line 10) | def __init__(self):
    method forward (line 15) | def forward(self, x):
  class TestConstrainedConcretize (line 21) | class TestConstrainedConcretize(TestCase):
    method __init__ (line 22) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 25) | def test(self):

FILE: tests/test_conv.py
  class cnn_model (line 7) | class cnn_model(nn.Module):
    method __init__ (line 8) | def __init__(self, layers, padding, stride, linear):
    method forward (line 25) | def forward(self, x):
  class TestConv (line 29) | class TestConv(TestCase):
    method __init__ (line 30) | def __init__(self, methodName='runTest', generate=False,
    method test (line 37) | def test(self):

FILE: tests/test_conv1d.py
  class Model (line 12) | class Model(nn.Module):
    method __init__ (line 13) | def __init__(self, kernel_size=2, stride=1, padding=0, in_features=1,o...
    method forward (line 27) | def forward(self, *inputs,debug=False):
  class TestConv1D (line 45) | class TestConv1D(TestCase):
    method __init__ (line 46) | def __init__(self, methodName='runTest', generate=False,
    method test (line 53) | def test(self):

FILE: tests/test_distinct_patches.py
  function reset_seed (line 14) | def reset_seed(seed=1234):
  class cnn_4layer_b (line 21) | class cnn_4layer_b(nn.Module):
    method __init__ (line 22) | def __init__(self, paddingA, paddingB):
    method forward (line 33) | def forward(self, x):
  class TestDistinctPatches (line 45) | class TestDistinctPatches(TestCase):
    method __init__ (line 46) | def __init__(self, methodName='runTest', generate=False,
    method run_conv_mode (line 69) | def run_conv_mode(self, model, img, conv_mode):
    method test (line 83) | def test(self):

FILE: tests/test_examples.py
  function download_data_language (line 22) | def download_data_language():
  function test_transformer (line 30) | def test_transformer():
  function test_lstm (line 39) | def test_lstm():
  function test_lstm_seq (line 50) | def test_lstm_seq():
  function test_simple_verification (line 57) | def test_simple_verification():
  function test_custom_op (line 63) | def test_custom_op():
  function test_efficient_convolution (line 69) | def test_efficient_convolution():
  function test_two_node (line 75) | def test_two_node():
  function test_simple_training (line 81) | def test_simple_training():
  function test_cifar_training (line 88) | def test_cifar_training():
  function test_weight_perturbation (line 96) | def test_weight_perturbation():
  function test_tinyimagenet (line 104) | def test_tinyimagenet():
  function test_imagenet (line 116) | def test_imagenet():
  function test_release (line 129) | def test_release():

FILE: tests/test_examples_ci.py
  function custom_run (line 9) | def custom_run(*args, **kwargs):
  function run_tests (line 17) | def run_tests():

FILE: tests/test_general_nonlinear.py
  class Sin (line 16) | class Sin(nn.Module):
    method forward (line 17) | def forward(self, x):
  function cifar_model_wide (line 21) | def cifar_model_wide():
  function bab (line 36) | def bab(model_ori, data, target, norm, eps, data_max=None, data_min=None...
  function test (line 92) | def test(device=DEFAULT_DEVICE, dtype=DEFAULT_DTYPE):

FILE: tests/test_general_shape.py
  class GeneralShapeModel (line 12) | class GeneralShapeModel(nn.Module):
    method __init__ (line 13) | def __init__(self):
    method forward (line 23) | def forward(self, x, w):
  class TestGeneralShape (line 56) | class TestGeneralShape(TestCase):
    method __init__ (line 57) | def __init__(self, methodName='runTest', seed=1, generate=False,
    method test (line 62) | def test(self):

FILE: tests/test_identity.py
  class TestIdentity (line 8) | class TestIdentity(TestCase):
    method __init__ (line 9) | def __init__(self, methodName='runTest', device=DEFAULT_DEVICE, dtype=...
    method test (line 12) | def test(self):

FILE: tests/test_invprop.py
  class SimpleExampleModel (line 12) | class SimpleExampleModel(nn.Module):
    method __init__ (line 13) | def __init__(self):
    method forward (line 19) | def forward(self, x):
  class TestInvpropSimpleExample (line 28) | class TestInvpropSimpleExample(TestCase):
    method __init__ (line 29) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 34) | def test(self):
  class TestInvpropOODExample (line 100) | class TestInvpropOODExample(TestCase):
    method __init__ (line 102) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 106) | def test(self):

FILE: tests/test_jacobian.py
  class TestJacobian (line 15) | class TestJacobian(TestCase):
    method __init__ (line 16) | def __init__(self, methodName='runTest', generate=False,
    method test (line 23) | def test(self):
    method test_concat_jacobian (line 41) | def test_concat_jacobian(self):

FILE: tests/test_language_models.py
  function prepare_data (line 15) | def prepare_data():
  function train (line 41) | def train():
  function read_res (line 57) | def read_res():
  function evaluate (line 61) | def evaluate():
  function gen_ref (line 83) | def gen_ref():
  function check (line 91) | def check():
  function test (line 102) | def test():

FILE: tests/test_linear_cnn_model.py
  class LinearCNNModel (line 12) | class LinearCNNModel(nn.Module):
    method __init__ (line 13) | def __init__(self):
    method forward (line 17) | def forward(self, x):
  class TestLinearCNNModel (line 22) | class TestLinearCNNModel(TestLinearModel):
    method __init__ (line 23) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method compute_and_compare_bounds (line 27) | def compute_and_compare_bounds(self, eps, norm, IBP, method):

FILE: tests/test_linear_model.py
  class LinearModel (line 10) | class LinearModel(nn.Module):
    method __init__ (line 11) | def __init__(self):
    method forward (line 15) | def forward(self, x):
  class TestLinearModel (line 19) | class TestLinearModel(TestCase):
    method __init__ (line 20) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method compute_and_compare_bounds (line 24) | def compute_and_compare_bounds(self, eps, norm, IBP, method):
    method test_Linf_forward (line 54) | def test_Linf_forward(self):
    method test_Linf_backward (line 58) | def test_Linf_backward(self):
    method test_Linf_IBP (line 62) | def test_Linf_IBP(self):
    method test_Linf_backward_IBP (line 66) | def test_Linf_backward_IBP(self):
    method test_L2_forward (line 70) | def test_L2_forward(self):
    method test_L2_backward (line 74) | def test_L2_backward(self):
    method test_L2_IBP (line 78) | def test_L2_IBP(self):
    method test_L2_backward_IBP (line 82) | def test_L2_backward_IBP(self):
    method test_L1_forward (line 86) | def test_L1_forward(self):
    method test_L1_backward (line 90) | def test_L1_backward(self):
    method test_L1_IBP (line 94) | def test_L1_IBP(self):
    method test_L1_backward_IBP (line 98) | def test_L1_backward_IBP(self):

FILE: tests/test_maxpool.py
  class Model (line 11) | class Model(nn.Module):
    method __init__ (line 12) | def __init__(self, kernel_size=4, stride=4, padding=0, conv_padding=0):
    method forward (line 24) | def forward(self, *inputs):
  class TestMaxPool (line 44) | class TestMaxPool(TestCase):
    method __init__ (line 45) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 51) | def test(self):

FILE: tests/test_min_max.py
  class Test_Model (line 10) | class Test_Model(nn.Module):
    method __init__ (line 11) | def __init__(self):
    method forward (line 35) | def forward(self, x):
  class TestMinMax (line 38) | class TestMinMax(TestCase):
    method __init__ (line 39) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 44) | def test(self):

FILE: tests/test_perturbation.py
  class ToyModel (line 16) | class ToyModel(nn.Module):
    method __init__ (line 18) | def __init__(self):
    method forward (line 24) | def forward(self, x):
  class TestPerturbation (line 31) | class TestPerturbation(TestCase):
    method __init__ (line 37) | def __init__(self, methodName='runTest', seed=1, generate=False,
    method test (line 41) | def test(self):

FILE: tests/test_rectangle_patches.py
  class cnn_4layer_resnet (line 12) | class cnn_4layer_resnet(nn.Module):
    method __init__ (line 13) | def __init__(self):
    method forward (line 21) | def forward(self, x):
  class TestResnetPatches (line 32) | class TestResnetPatches(TestCase):
    method __init__ (line 33) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 39) | def test(self):

FILE: tests/test_resnet_patches.py
  class TestResnetPatches (line 13) | class TestResnetPatches(TestCase):
    method __init__ (line 14) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 20) | def test(self):

FILE: tests/test_s_shaped.py
  class test_model (line 9) | class test_model(nn.Module):
    method __init__ (line 10) | def __init__(self, act_func):
    method forward (line 14) | def forward(self, x):
  function sigmoid (line 17) | def sigmoid(x):
  function sin (line 20) | def sin(x):
  function verify_bounds (line 24) | def verify_bounds(model, input_lb, input_ub, lb, ub):
  class TestSShaped (line 50) | class TestSShaped(TestCase):
    method __init__ (line 51) | def __init__(self, methodName='runTest', generate=False,
    method _run_bound_test (line 58) | def _run_bound_test(self, model, input_lb, input_ub, methods):
    method test (line 70) | def test(self):

FILE: tests/test_save_intermediate.py
  class test_model (line 7) | class test_model(nn.Module):
    method __init__ (line 8) | def __init__(self):
    method forward (line 21) | def forward(self, x):
  class TestSave (line 25) | class TestSave(TestCase):
    method __init__ (line 26) | def __init__(self, methodName='runTest', device=DEFAULT_DEVICE, dtype=...
    method test (line 29) | def test(self, gen_ref=False):

FILE: tests/test_simple_verification.py
  function mnist_model (line 11) | def mnist_model():
  class TestSimpleVerification (line 24) | class TestSimpleVerification(TestCase):
    method __init__ (line 25) | def __init__(self, methodName='runTest', device=DEFAULT_DEVICE, dtype=...
    method test (line 28) | def test(self):

FILE: tests/test_state_dict_name.py
  class FeatureExtraction (line 7) | class FeatureExtraction(nn.Module):
    method __init__ (line 8) | def __init__(self):
    method forward (line 14) | def forward(self, x):
  class cnn_MNIST (line 22) | class cnn_MNIST(nn.Module):
    method __init__ (line 23) | def __init__(self):
    method forward (line 28) | def forward(self, x):
  class TestStateDictName (line 33) | class TestStateDictName(TestCase):
    method __init__ (line 34) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 37) | def test(self):

FILE: tests/test_tensor_storage.py
  class TestTensorStorage (line 9) | class TestTensorStorage(TestCase):
    method __init__ (line 10) | def __init__(self, methodName='runTest', device=DEFAULT_DEVICE, dtype=...
    method test_content (line 13) | def test_content(self, seed=123):
    method test_tensor_call (line 56) | def test_tensor_call(self, seed=123):
    method test_size_queue (line 86) | def test_size_queue(self):
    method test_size_stack (line 139) | def test_size_stack(self):

FILE: tests/test_upsample.py
  class Model (line 9) | class Model(nn.Module):
    method __init__ (line 11) | def __init__(self,
    method forward (line 34) | def forward(self, input_z):
  class ModelReducedCGAN (line 45) | class ModelReducedCGAN(nn.Module):
    method __init__ (line 46) | def __init__(self):
    method forward (line 73) | def forward(self, input_z):
  function recursive_allclose (line 101) | def recursive_allclose(a, b: dict, verbose=False, prefix=''):
  class TestUpSample (line 131) | class TestUpSample(TestCase):
    method __init__ (line 132) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 137) | def test(self, seed=123):
    method test_instance (line 144) | def test_instance(self, kernel_size=3, scaling_factor=2, stride=1, pad...
  class TestReducedCGAN (line 176) | class TestReducedCGAN(TestCase):
    method __init__ (line 178) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test (line 183) | def test(self, seed=456):

FILE: tests/test_vision_models.py
  class cnn_4layer_test (line 8) | class cnn_4layer_test(nn.Module):
    method __init__ (line 9) | def __init__(self):
    method forward (line 17) | def forward(self, x):
  class TestVisionModels (line 27) | class TestVisionModels(TestCase):
    method __init__ (line 28) | def __init__(self, methodName='runTest', ref_name='vision_test_data', ...
    method setUp (line 35) | def setUp(self):
    method verify_bounds (line 44) | def verify_bounds(self, model, x, IBP, method, forward_ret, lb_name, u...
    method test_bounds (line 73) | def test_bounds(self, bound_opts=None, optimize = True):

FILE: tests/test_vision_models_hardtanh.py
  class cnn_4layer_test_hardtanh (line 7) | class cnn_4layer_test_hardtanh(nn.Module):
    method __init__ (line 8) | def __init__(self, in_ch, in_dim, width=2, linear_size=256):
    method forward (line 15) | def forward(self, x):
  class TestCustomVisionModel (line 24) | class TestCustomVisionModel(TestVisionModels):
    method __init__ (line 25) | def __init__(self, methodName='runTest', model=cnn_4layer_test_hardtan...
    method test_bounds (line 28) | def test_bounds(self, bound_opts=None, optimize=False):

FILE: tests/test_weight_perturbation.py
  class TestWeightPerturbation (line 12) | class TestWeightPerturbation(TestCase):
    method __init__ (line 13) | def __init__(self, methodName='runTest', generate=False, device=DEFAUL...
    method test_training (line 20) | def test_training(self):
    method verify_bounds (line 33) | def verify_bounds(self, model, x, IBP, method, forward_ret, lb_name, u...
    method test_perturbation (line 55) | def test_perturbation(self):

FILE: tests/testcase.py
  class TestCase (line 9) | class TestCase(unittest.TestCase):
    method __init__ (line 12) | def __init__(self, methodName='runTest', seed=1, ref_name=None, genera...
    method set_seed (line 30) | def set_seed(self, seed):
    method setUp (line 36) | def setUp(self):
    method save (line 43) | def save(self):
    method check (line 48) | def check(self):
    method _assert_equal (line 66) | def _assert_equal(self, a, b):
    method _assert_array_equal (line 74) | def _assert_array_equal(self, a, b, msg=None):
    method _assert_tensor_equal (line 84) | def _assert_tensor_equal(self, a, b, msg=None):
  function _to (line 95) | def _to(obj, device=None, dtype=None, inplace=False):
  function set_default_dtype_device (line 118) | def set_default_dtype_device(dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE):
Condensed preview — 249 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (1,612K chars).
[
  {
    "path": ".github/ISSUE_TEMPLATE/bug_report.md",
    "chars": 1469,
    "preview": "---\nname: Bug report\nabout: Create a report to help us improve\ntitle: ''\nlabels: ''\nassignees: ''\n\n---\n\n**Describe the b"
  },
  {
    "path": ".gitignore",
    "chars": 262,
    "preview": "tmp\nbuild\n__pycache__\n*.egg-info\ndist\n*.swp\n*.swo\n*.log\n.trace_graph\nVerified_ret*.npy\nVerified-acc*.npy\nvnn-comp_*.npz\n"
  },
  {
    "path": ".readthedocs.yaml",
    "chars": 510,
    "preview": "# Read the Docs configuration file\n# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details\n\n# Requir"
  },
  {
    "path": "CONTRIBUTORS",
    "chars": 1928,
    "preview": "Team leaders:\n* Faculty: Huan Zhang (huan@huan-zhang.com), UIUC\n* Student: Xiangru Zhong (xiangru4@illinois.edu), UIUC\n\n"
  },
  {
    "path": "LICENSE",
    "chars": 2326,
    "preview": "Copyright (C) 2021-2025 The α,β-CROWN Team\nSee CONTRIBUTORS for the list of all contributors and their affiliations.\n   "
  },
  {
    "path": "README.md",
    "chars": 18792,
    "preview": "# auto_LiRPA: Automatic Linear Relaxation based Perturbation Analysis for Neural Networks\n\n[![Documentation Status](http"
  },
  {
    "path": "auto_LiRPA/__init__.py",
    "chars": 1581,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/backward_bound.py",
    "chars": 48867,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/beta_crown.py",
    "chars": 11699,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/bound_general.py",
    "chars": 71966,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/bound_multi_gpu.py",
    "chars": 7984,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/bound_op_map.py",
    "chars": 1826,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/bound_ops.py",
    "chars": 1209,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/bounded_tensor.py",
    "chars": 4662,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/concretize_bounds.py",
    "chars": 18411,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/concretize_func.py",
    "chars": 37368,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/cuda/cuda_kernels.cu",
    "chars": 1234,
    "preview": "#include <torch/extension.h>\n\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n#include <vector>\n\n__global__ void cuda_doubl"
  },
  {
    "path": "auto_LiRPA/cuda/cuda_utils.cpp",
    "chars": 981,
    "preview": "#include <torch/extension.h>\n\n#include <vector>\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x \" must be a CU"
  },
  {
    "path": "auto_LiRPA/cuda_utils.py",
    "chars": 5720,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/edit_graph.py",
    "chars": 3528,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/eps_scheduler.py",
    "chars": 12824,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/forward_bound.py",
    "chars": 13944,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/interval_bound.py",
    "chars": 9728,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/jacobian.py",
    "chars": 9102,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/linear_bound.py",
    "chars": 2374,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/__init__.py",
    "chars": 1946,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/activation_base.py",
    "chars": 16042,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/activations.py",
    "chars": 21516,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/add_sub.py",
    "chars": 10247,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/base.py",
    "chars": 28212,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/bivariate.py",
    "chars": 22958,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/clampmult.py",
    "chars": 9588,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/constant.py",
    "chars": 6885,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/convex_concave.py",
    "chars": 14277,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/convolution.py",
    "chars": 51078,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/cut_ops.py",
    "chars": 38587,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/dropout.py",
    "chars": 3868,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/dtype.py",
    "chars": 3089,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/gelu.py",
    "chars": 20634,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/indexing.py",
    "chars": 12591,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/jacobian.py",
    "chars": 2915,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/leaf.py",
    "chars": 4370,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/linear.py",
    "chars": 56616,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/logical.py",
    "chars": 2409,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/minmax.py",
    "chars": 14358,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/normalization.py",
    "chars": 16096,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/pooling.py",
    "chars": 38281,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/reduce.py",
    "chars": 11584,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/relu.py",
    "chars": 63361,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/reshape.py",
    "chars": 17390,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/resize.py",
    "chars": 20599,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/rnn.py",
    "chars": 3525,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/s_shaped.py",
    "chars": 57026,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/shape.py",
    "chars": 2140,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/slice_concat.py",
    "chars": 20703,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/softmax.py",
    "chars": 2850,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/solver_utils.py",
    "chars": 1560,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/tile.py",
    "chars": 3232,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/operators/trigonometric.py",
    "chars": 25848,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/opt_pruner.py",
    "chars": 14726,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/optimize_graph.py",
    "chars": 16218,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/optimized_bounds.py",
    "chars": 49517,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/output_constraints.py",
    "chars": 16839,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/parse_graph.py",
    "chars": 10633,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/patches.py",
    "chars": 36428,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/perturbations.py",
    "chars": 36482,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/solver_module.py",
    "chars": 10240,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/tools.py",
    "chars": 5535,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/utils.py",
    "chars": 19152,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "auto_LiRPA/wrapper.py",
    "chars": 1924,
    "preview": "#########################################################################\n##   This file is part of the auto_LiRPA libra"
  },
  {
    "path": "doc/.gitignore",
    "chars": 41,
    "preview": "_build\nsections\n*.md\n!src/*.md\n!README.md"
  },
  {
    "path": "doc/Makefile",
    "chars": 634,
    "preview": "# Minimal makefile for Sphinx documentation\n#\n\n# You can set these variables from the command line, and also\n# from the "
  },
  {
    "path": "doc/README.md",
    "chars": 492,
    "preview": "# Documentation\n\nThis directory contains source files for building our documentation.\nPlease view the compiled documenta"
  },
  {
    "path": "doc/api.rst",
    "chars": 850,
    "preview": "API Usage\n======================================\n\n.. autoclass:: auto_LiRPA.BoundedModule\n\n   .. autofunction:: auto_LiR"
  },
  {
    "path": "doc/conf.py",
    "chars": 2960,
    "preview": "# Configuration file for the Sphinx documentation builder.\n#\n# This file only contains a selection of the most common op"
  },
  {
    "path": "doc/index.rst",
    "chars": 1163,
    "preview": ".. auto_LiRPA documentation master file, created by\n   sphinx-quickstart on Wed Jul 14 21:56:10 2021.\n   You can adapt t"
  },
  {
    "path": "doc/process.py",
    "chars": 2363,
    "preview": "\"\"\" Process source files before running Sphinx\"\"\"\nimport re\nimport os\nimport shutil\nfrom pygit2 import Repository\n\nrepo "
  },
  {
    "path": "examples/.gitignore",
    "chars": 11,
    "preview": "auto_LiRPA\n"
  },
  {
    "path": "examples/__init__.py",
    "chars": 0,
    "preview": ""
  },
  {
    "path": "examples/language/.gitignore",
    "chars": 70,
    "preview": "model*\n!modeling*\nlog*\nres_test.pkl\nckpt_*\ndata_language.tar.gz\ndata/\n"
  },
  {
    "path": "examples/language/Transformer/Transformer.py",
    "chars": 4922,
    "preview": "# coding=utf-8\n# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.\n# Copyright (c) 2018,"
  },
  {
    "path": "examples/language/Transformer/__init__.py",
    "chars": 0,
    "preview": ""
  },
  {
    "path": "examples/language/Transformer/modeling.py",
    "chars": 10464,
    "preview": "# coding=utf-8\n# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.\n# Copyright (c) 2018,"
  },
  {
    "path": "examples/language/Transformer/utils.py",
    "chars": 2541,
    "preview": "# coding=utf-8\n# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.\n# Copyright (c) 2018,"
  },
  {
    "path": "examples/language/data_utils.py",
    "chars": 783,
    "preview": "import random\nimport json\nfrom auto_LiRPA.utils import logger\n\ndef load_data_sst():\n    data = []\n    for split in ['tra"
  },
  {
    "path": "examples/language/language_utils.py",
    "chars": 1427,
    "preview": "from auto_LiRPA.utils import logger\nimport numpy as np\n\ndef build_vocab(data_train, min_word_freq, dump=False, include=["
  },
  {
    "path": "examples/language/lstm.py",
    "chars": 5331,
    "preview": "import os\nimport shutil\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom auto_LiRPA.utils import "
  },
  {
    "path": "examples/language/oracle.py",
    "chars": 2218,
    "preview": "import torch\nfrom auto_LiRPA.utils import logger\nfrom auto_LiRPA import PerturbationSynonym\nfrom data_utils import get_b"
  },
  {
    "path": "examples/language/preprocess/pre_compute_lm_scores.py",
    "chars": 2547,
    "preview": "# Ref: https://worksheets.codalab.org/rest/bundles/0x3f614472f4a14393b3d85d5568114591/contents/blob/precompute_lm_scores"
  },
  {
    "path": "examples/language/preprocess/preprocess_sst.py",
    "chars": 3213,
    "preview": "import random, json\n\ndef load_data_sst():\n    # training data\n    path = \"train-nodes.tsv\"\n    data_train_all_nodes = []"
  },
  {
    "path": "examples/language/train.py",
    "chars": 14314,
    "preview": "import argparse\nimport random\nimport pickle\nimport os\nimport pdb\nimport time\nimport logging\nimport numpy as np\nimport to"
  },
  {
    "path": "examples/sequence/.gitignore",
    "chars": 12,
    "preview": "model/\ndata/"
  },
  {
    "path": "examples/sequence/__init__.py",
    "chars": 0,
    "preview": ""
  },
  {
    "path": "examples/sequence/data_utils.py",
    "chars": 771,
    "preview": "import random\nfrom torchvision import transforms\nfrom torchvision.datasets.mnist import MNIST as mnist\n\ndef load_data():"
  },
  {
    "path": "examples/sequence/lstm.py",
    "chars": 2669,
    "preview": "import os\nimport shutil\nimport torch\nimport torch.nn as nn\nfrom auto_LiRPA.utils import logger\n\nclass LSTMCore(nn.Module"
  },
  {
    "path": "examples/sequence/train.py",
    "chars": 5352,
    "preview": "import argparse\nimport random\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nfrom"
  },
  {
    "path": "examples/simple/invprop.py",
    "chars": 4761,
    "preview": "\"\"\"\nA toy example for bounding neural network outputs under input perturbations using INVPROP\n\nSee https://arxiv.org/abs"
  },
  {
    "path": "examples/simple/lp_full.py",
    "chars": 11428,
    "preview": "\"\"\"\nA simple example for bounding neural network outputs using LP/MIP solvers.\n\nAuto_LiRPA supports constructing LP/MIP "
  },
  {
    "path": "examples/simple/mip_lp_solver.py",
    "chars": 5070,
    "preview": "\"\"\"\nA simple example for bounding neural network outputs using LP/MIP solvers.\n\nAuto_LiRPA supports constructing LP/MIP "
  },
  {
    "path": "examples/simple/toy.py",
    "chars": 2475,
    "preview": "\"\"\"\nA toy example for bounding neural network outputs under input perturbations.\n\"\"\"\nimport torch\nfrom collections impor"
  },
  {
    "path": "examples/vision/.gitignore",
    "chars": 69,
    "preview": "exp\nexp_inv\n__pycache__\nmodel_*\n!model_gurobi.py\nsaved_models\nconfig\n"
  },
  {
    "path": "examples/vision/bound_option.py",
    "chars": 6092,
    "preview": "\"\"\"\nA simple example for bounding neural network outputs with different bound options on ReLU activation functions.\n\n\"\"\""
  },
  {
    "path": "examples/vision/cifar_training.py",
    "chars": 16695,
    "preview": "import argparse\nimport multiprocessing\nimport random\nimport time\nimport logging\nimport os\n\nimport torch.optim as optim\ni"
  },
  {
    "path": "examples/vision/custom_op.py",
    "chars": 5127,
    "preview": "\"\"\" A example for custom operators.\n\nIn this example, we create a custom operator called \"PlusConstant\", which can\nbe wr"
  },
  {
    "path": "examples/vision/data/.gitignore",
    "chars": 12,
    "preview": "MNIST\ncifar*"
  },
  {
    "path": "examples/vision/data/ImageNet64/imagenet_data_loader.py",
    "chars": 1512,
    "preview": "import os\n\nimport numpy as np\nfrom PIL import Image\n\n\nclass DatasetDownsampledImageNet():\n    def __init__(self):\n      "
  },
  {
    "path": "examples/vision/data/tinyImageNet/.gitignore",
    "chars": 19,
    "preview": "tiny-imagenet-200*\n"
  },
  {
    "path": "examples/vision/data/tinyImageNet/tinyimagenet_download.sh",
    "chars": 764,
    "preview": "#!/bin/bash\n\n# download and unzip dataset\nwget http://cs231n.stanford.edu/tiny-imagenet-200.zip\nunzip tiny-imagenet-200."
  },
  {
    "path": "examples/vision/datasets.py",
    "chars": 8851,
    "preview": "import multiprocessing\nimport torch\nfrom torch.utils import data\nfrom functools import partial\nimport torchvision.transf"
  },
  {
    "path": "examples/vision/efficient_convolution.py",
    "chars": 3361,
    "preview": "\"\"\"\nDemonstration of efficient convolutional network implementation in auto_LiRPA.\n\nauto_LiRPA library supports an effic"
  },
  {
    "path": "examples/vision/imagenet_training.py",
    "chars": 15903,
    "preview": "import random\nimport time\nimport argparse\nimport multiprocessing\nimport logging\nimport torch.optim as optim\nfrom torch.n"
  },
  {
    "path": "examples/vision/jacobian.py",
    "chars": 5417,
    "preview": "\"\"\"Examples of computing Jacobian bounds.\n\nWe show examples of:\n- Computing Jacobian bounds\n- Computing Linf local Lipsc"
  },
  {
    "path": "examples/vision/models/__init__.py",
    "chars": 1602,
    "preview": "from models.resnet import model_resnet\nfrom models.feedforward import *\nfrom models.resnext import *\nfrom models.resnext"
  },
  {
    "path": "examples/vision/models/densenet.py",
    "chars": 3546,
    "preview": "'''DenseNet in PyTorch.\nhttps://github.com/kuangliu/pytorch-cifar\n'''\n\n\nimport math\nimport torch\nimport torch.nn as nn\ni"
  },
  {
    "path": "examples/vision/models/densenet_imagenet.py",
    "chars": 3554,
    "preview": "'''DenseNet in PyTorch.\nhttps://github.com/kuangliu/pytorch-cifar\n'''\n\n\nimport math\nimport torch\nimport torch.nn as nn\ni"
  },
  {
    "path": "examples/vision/models/densenet_no_bn.py",
    "chars": 3392,
    "preview": "'''DenseNet in PyTorch.\nhttps://github.com/kuangliu/pytorch-cifar\n'''\n\n\nimport math\nimport torch\nimport torch.nn as nn\ni"
  },
  {
    "path": "examples/vision/models/feedforward.py",
    "chars": 6429,
    "preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom auto_LiRPA import PerturbationLpNorm, BoundedPar"
  },
  {
    "path": "examples/vision/models/mobilenet.py",
    "chars": 3063,
    "preview": "'''MobileNetV2 in PyTorch.\n\nSee the paper \"Inverted Residuals and Linear Bottlenecks:\nMobile Networks for Classification"
  },
  {
    "path": "examples/vision/models/resnet.py",
    "chars": 2878,
    "preview": "'''\nResNet used in https://arxiv.org/pdf/1805.12514.pdf\nhttps://github.com/locuslab/convex_adversarial/blob/0d11e671ad93"
  },
  {
    "path": "examples/vision/models/resnet18.py",
    "chars": 4257,
    "preview": "'''ResNet in PyTorch.\n\nFor Pre-activation ResNet, see 'preact_resnet.py'.\n\nReference:\n[1] Kaiming He, Xiangyu Zhang, Sha"
  },
  {
    "path": "examples/vision/models/resnext.py",
    "chars": 3906,
    "preview": "'''ResNeXt in PyTorch.\nSee the paper \"Aggregated Residual Transformations for Deep Neural Networks\" for more details.\nht"
  },
  {
    "path": "examples/vision/models/resnext_imagenet64.py",
    "chars": 3507,
    "preview": "'''ResNeXt in PyTorch.\nSee the paper \"Aggregated Residual Transformations for Deep Neural Networks\" for more details.\nht"
  },
  {
    "path": "examples/vision/models/vnncomp_resnet.py",
    "chars": 6472,
    "preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass BasicBlock(nn.Module):\n    expansion = 1\n\n   "
  },
  {
    "path": "examples/vision/models/wide_resnet_cifar.py",
    "chars": 4470,
    "preview": "import torch\nimport torch.nn as nn\nimport torch.nn.init as init\nimport torch.nn.functional as F\nfrom torch.autograd impo"
  },
  {
    "path": "examples/vision/models/wide_resnet_imagenet64.py",
    "chars": 3669,
    "preview": "import torch\nimport torch.nn as nn\nimport torch.nn.init as init\nimport torch.nn.functional as F\n\nimport sys\nimport numpy"
  },
  {
    "path": "examples/vision/save_intermediate_bound.py",
    "chars": 1959,
    "preview": "\"\"\"\nA simple example for saving intermediate bounds.\n\"\"\"\nimport os\nimport torch\nimport torch.nn as nn\nimport torchvision"
  },
  {
    "path": "examples/vision/simple_training.py",
    "chars": 12066,
    "preview": "\"\"\"\nA simple script to train certified defense using the auto_LiRPA library.\n\nWe compute output bounds under input pertu"
  },
  {
    "path": "examples/vision/simple_verification.py",
    "chars": 8647,
    "preview": "\"\"\"\nA simple example for bounding neural network outputs under input perturbations.\n\nThis example serves as a skeleton f"
  },
  {
    "path": "examples/vision/tinyimagenet_training.py",
    "chars": 15991,
    "preview": "import os\nimport random\nimport time\nimport argparse\nimport multiprocessing\nimport logging\nimport torch.optim as optim\nfr"
  },
  {
    "path": "examples/vision/verify_two_node.py",
    "chars": 2956,
    "preview": "\"\"\"\nExample for multi-node perturbation. An input image is splited to two parts\nwhere each part is perturbed respectivel"
  },
  {
    "path": "examples/vision/weight_perturbation_training.py",
    "chars": 15544,
    "preview": "\"\"\"\nA simple example for certified robustness against model weight perturbations.\n\nSince our framework works on general "
  },
  {
    "path": "setup.py",
    "chars": 2251,
    "preview": "from setuptools import setup, find_packages\nfrom pathlib import Path\n\n# Check PyTorch version\npytorch_version_l = '2.0.0"
  },
  {
    "path": "tests/.gitignore",
    "chars": 7,
    "preview": ".cache\n"
  },
  {
    "path": "tests/data/.gitignore",
    "chars": 48,
    "preview": "cifar-10-python.tar.gz\ncifar-10-batches-py\nMNIST"
  },
  {
    "path": "tests/test_1d_activation.py",
    "chars": 7917,
    "preview": "\"\"\"Test one dimensional activation functions (e.g., ReLU, tanh, exp, sin, etc)\"\"\"\nimport functools\n\nimport pytest\nimport"
  },
  {
    "path": "tests/test_2d_activation.py",
    "chars": 6203,
    "preview": "\"\"\"Test two dimensional activation functions (e.g., min, max, etc)\"\"\"\nimport tqdm\nimport torch\nimport torch.nn as nn\nfro"
  },
  {
    "path": "tests/test_avgpool.py",
    "chars": 5149,
    "preview": "import torch\nimport torch.nn as nn\nimport numpy as np\nfrom auto_LiRPA import BoundedModule, BoundedTensor\nfrom auto_LiRP"
  },
  {
    "path": "tests/test_bound_ops.py",
    "chars": 6270,
    "preview": "\"\"\"Test classes for bound operators\"\"\"\nimport torch\nfrom auto_LiRPA.bound_ops import *\nfrom auto_LiRPA.linear_bound impo"
  },
  {
    "path": "tests/test_branching_heuristics.py",
    "chars": 3906,
    "preview": "import sys\nimport torch\nfrom types import SimpleNamespace\n\nsys.path.insert(0, '../complete_verifier')\n\nfrom heuristics.b"
  },
  {
    "path": "tests/test_clip_domains.py",
    "chars": 15709,
    "preview": "\"\"\"\nTests clip_domains\n\nTo run tests: py.test             test_clip_domains.py\n          or: python -m pytest    test_cl"
  },
  {
    "path": "tests/test_constant.py",
    "chars": 2417,
    "preview": "\"\"\"Test BoundConstant\"\"\"\nimport torch\nimport os\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision"
  },
  {
    "path": "tests/test_constrained_concretize.py",
    "chars": 2996,
    "preview": "\"\"\"Test optimized bounds in simple_verification.\"\"\"\nimport torch\nimport numpy as np\nfrom auto_LiRPA import BoundedModule"
  },
  {
    "path": "tests/test_conv.py",
    "chars": 3514,
    "preview": "import torch\nimport torch.nn as nn\nfrom auto_LiRPA import BoundedModule, BoundedTensor\nfrom auto_LiRPA.perturbations imp"
  },
  {
    "path": "tests/test_conv1d.py",
    "chars": 6473,
    "preview": "\"\"\"Test Conv1d.\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom auto_LiRPA import BoundedMo"
  },
  {
    "path": "tests/test_distinct_patches.py",
    "chars": 3937,
    "preview": "import torch\nimport random\nimport numpy as np\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nf"
  },
  {
    "path": "tests/test_examples.py",
    "chars": 4786,
    "preview": "\"\"\"Test all the examples before release.\n\nThis script is expected be manually run and is not used in automatic tests.\"\"\""
  },
  {
    "path": "tests/test_examples_ci.py",
    "chars": 1188,
    "preview": "import subprocess\nimport traceback\n\nimport test_examples\n\noriginal_subprocess_run = subprocess.run\n\n\ndef custom_run(*arg"
  },
  {
    "path": "tests/test_general_nonlinear.py",
    "chars": 3596,
    "preview": "import sys\nimport pytest\nimport torch.nn as nn\n\nsys.path.insert(0, '../complete_verifier')\n\nimport arguments\nfrom beta_C"
  },
  {
    "path": "tests/test_general_shape.py",
    "chars": 3601,
    "preview": "\"\"\" Test inputs of general shapes (especially for matmul)\"\"\"\nimport torch\nimport torch.nn as nn\nimport numpy as np\n\nfrom"
  },
  {
    "path": "tests/test_identity.py",
    "chars": 1039,
    "preview": "\"\"\"Test a model with an nn.Identity layer only\"\"\"\nimport torch\nimport torch.nn as nn\nfrom auto_LiRPA import BoundedModul"
  },
  {
    "path": "tests/test_invprop.py",
    "chars": 7449,
    "preview": "\"\"\"Test INVPROP.\"\"\"\nimport sys\nsys.path.append('../complete_verifier')\nfrom complete_verifier.load_model import unzip_an"
  },
  {
    "path": "tests/test_jacobian.py",
    "chars": 2152,
    "preview": "# pylint: disable=wrong-import-position\n\"\"\"Test Jacobian bounds.\"\"\"\nimport sys\nimport torch\nimport torch.nn as nn\n\nsys.p"
  },
  {
    "path": "tests/test_language_models.py",
    "chars": 4615,
    "preview": "\"\"\"Test classes for Transformer and LSTM on language tasks\"\"\"\nimport os\nimport argparse\nimport pickle\nimport torch\nfrom "
  },
  {
    "path": "tests/test_linear_cnn_model.py",
    "chars": 2733,
    "preview": "\"\"\"Test bounds on a 1 layer CNN network.\"\"\"\nimport torch.nn as nn\nfrom auto_LiRPA import BoundedModule, BoundedTensor\nfr"
  },
  {
    "path": "tests/test_linear_model.py",
    "chars": 3947,
    "preview": "\"\"\"Test bounds on a 1 layer linear network.\"\"\"\nimport torch.nn as nn\nfrom auto_LiRPA import BoundedModule, BoundedTensor"
  },
  {
    "path": "tests/test_maxpool.py",
    "chars": 4937,
    "preview": "\"\"\"Test max pooling.\"\"\"\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom auto_LiRPA import Bound"
  },
  {
    "path": "tests/test_min_max.py",
    "chars": 3070,
    "preview": "import os\nimport torch\nimport torch.nn as nn\nimport torchvision\nfrom auto_LiRPA import BoundedModule, BoundedTensor\nfrom"
  },
  {
    "path": "tests/test_perturbation.py",
    "chars": 4079,
    "preview": "\"\"\" Test different Perturbation classes\"\"\"\nimport torch\nimport torch.nn as nn\nimport numpy as np\n\nfrom auto_LiRPA import"
  },
  {
    "path": "tests/test_rectangle_patches.py",
    "chars": 2947,
    "preview": "import sys\nimport torch\nimport numpy as np\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nfrom"
  },
  {
    "path": "tests/test_resnet_patches.py",
    "chars": 2105,
    "preview": "import sys\nimport torch\nimport numpy as np\nimport torchvision\nimport models\nfrom auto_LiRPA import BoundedModule, Bounde"
  },
  {
    "path": "tests/test_s_shaped.py",
    "chars": 3818,
    "preview": "# pylint: disable=wrong-import-position\n\"\"\"Test S-shaped activation functions.\"\"\"\nimport torch\nimport torch.nn as nn\nfro"
  },
  {
    "path": "tests/test_save_intermediate.py",
    "chars": 2264,
    "preview": "import torch\nimport torch.nn as nn\nfrom auto_LiRPA import BoundedModule, BoundedTensor\nfrom auto_LiRPA.perturbations imp"
  },
  {
    "path": "tests/test_simple_verification.py",
    "chars": 2029,
    "preview": "\"\"\"Test optimized bounds in simple_verification.\"\"\"\nimport torch\nimport torch.nn as nn\nimport torchvision\nfrom auto_LiRP"
  },
  {
    "path": "tests/test_state_dict_name.py",
    "chars": 1766,
    "preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom auto_LiRPA import BoundedModule\nfrom testcase im"
  },
  {
    "path": "tests/test_tensor_storage.py",
    "chars": 9746,
    "preview": "\nimport random\nimport torch\nfrom complete_verifier.tensor_storage import StackTensorStorage, QueueTensorStorage\n\nfrom te"
  },
  {
    "path": "tests/test_upsample.py",
    "chars": 10273,
    "preview": "from collections import defaultdict\n\nfrom torch import nn\nfrom auto_LiRPA import BoundedModule, BoundedTensor\nfrom auto_"
  },
  {
    "path": "tests/test_vision_models.py",
    "chars": 7656,
    "preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom auto_LiRPA import BoundedModule, BoundedTensor\nf"
  },
  {
    "path": "tests/test_vision_models_hardtanh.py",
    "chars": 1473,
    "preview": "import torch.nn as nn\nimport torch.nn.functional as F\nfrom auto_LiRPA.perturbations import *\nfrom test_vision_models imp"
  },
  {
    "path": "tests/test_weight_perturbation.py",
    "chars": 5437,
    "preview": "import copy\nimport subprocess\nimport numpy as np\nfrom testcase import TestCase, DEFAULT_DEVICE, DEFAULT_DTYPE \nimport sy"
  },
  {
    "path": "tests/testcase.py",
    "chars": 4763,
    "preview": "import unittest\nimport random\nimport torch\nimport numpy as np\n\nDEFAULT_DEVICE = 'cpu'\nDEFAULT_DTYPE = torch.float32\n\ncla"
  }
]

// ... and 79 more files (download for full content)

About this extraction

This page contains the full source code of the KaidiXu/auto_LiRPA GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 249 files (1.5 MB), approximately 383.1k tokens, and a symbol index with 1697 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.

Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.

Copied to clipboard!