Resolve Failing Tests And Patch For Bounds in Linear Tree (#179)

This PR addresses several issues that have resulted in failing tests as
well as a patch for unscaled bounds in linear-tree. Here is everything
that is addressed in this PR

- Adding a free disk space step since runners are running out of space
for Github CI
- Update pyproject.toml to point to
https://github.com/All-Things-Optimal/linear-tree
- Update pyproject.toml to ignore ruff  "PLW1641" and "TRY003"
- Limit multiple big-M tranfsformation to one thread
- Skip MBM transformation tests until newest release of pyomo (>6.9.5)
due to bug that prevents mbm from transforming blocks
- Update ONNX parser to work with most recent version of torch for
convolutional neural networks
- Add tests to ONNX parser to check errors for strides and kernel_shape
attributes in onnx graph for convolution nodes and reshape nodes
- Skip import_networks.ipynb until tf2onnx addresses dependency issue
(last tf2onnx release was january 2024)
- Fix bug with unscaled input bounds being called before it is defined
in lt_formulations.py
- Address several ruff and ruff format issues

**Legal Acknowledgement**\
By contributing to this software project, I agree my contributions are
submitted under the BSD license.
I represent I am authorized to make the contributions and grant the
license.
If my employer has rights to intellectual property that includes these
contributions,
I represent that I have received permission to make contributions and
grant the required license on behalf of that employer.

---------

Co-authored-by: Ethan Sunshine <93541000+emsunshine@users.noreply.github.com>

Author: bammari

.github/workflows/tests.yml

@@ -10,6 +10,13 @@ jobs:
     runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v4
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          df -h
       - uses: actions/setup-python@v5
         with:
           python-version: "3.12"
@@ -23,6 +30,13 @@ jobs:
     runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v4
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          df -h
       - uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
@@ -33,6 +47,13 @@ jobs:
     runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v4
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          df -h
       - uses: actions/setup-python@v5
         with:
           python-version: "3.12"
@@ -46,6 +67,13 @@ jobs:
     runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v4
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          df -h
       - uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
@@ -74,6 +102,13 @@ jobs:
     runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v4
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          df -h
       - uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}

docs/notebooks/neuralnet/import_network.ipynb

@@ -864,7 +864,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -885,7 +885,7 @@
     "    torch.onnx.export(\n",
     "        model,\n",
     "        x,\n",
-    "        f,\n",
+    "        f.name,\n",
     "        input_names=[\"input\"],\n",
     "        output_names=[\"output\"],\n",
     "        dynamic_axes={\"input\": {0: \"batch_size\"}, \"output\": {0: \"batch_size\"}},\n",

docs/notebooks/neuralnet/mnist_example_convolutional.ipynb

@@ -362,7 +362,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -371,7 +371,7 @@
     "    torch.onnx.export(\n",
     "        model2,\n",
     "        x_temp,\n",
-    "        f,\n",
+    "        f.name,\n",
     "        input_names=[\"input\"],\n",
     "        output_names=[\"output\"],\n",
     "        dynamic_axes={\"input\": {0: \"batch_size\"}, \"output\": {0: \"batch_size\"}},\n",

docs/notebooks/neuralnet/mnist_example_dense.ipynb

@@ -154,14 +154,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 52,
+   "execution_count": null,
    "id": "c54d0dad",
    "metadata": {},
    "outputs": [],
    "source": [
-    "def write_onnx_to_file(onnx_model, path, file_name=\"output.onnx\"):\n",
-    "    from pathlib import Path\n",
+    "from pathlib import Path\n",
+    "\n",
     "\n",
+    "def write_onnx_to_file(onnx_model, path, file_name=\"output.onnx\"):\n",
     "    with open(Path(path) / file_name, \"wb\") as onnx_file:\n",
     "        onnx_file.write(onnx_model.SerializeToString())\n",
     "        print(f\"Onnx model written to {onnx_file.name}\")"
@@ -248,7 +249,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 55,
+   "execution_count": null,
    "id": "376f0893",
    "metadata": {
     "scrolled": false
@@ -268,14 +269,15 @@
     }
    ],
    "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
     "# defining initial data\n",
     "random.seed(10)\n",
     "data = generate_samples(5, f_bnds)\n",
     "\n",
     "\n",
     "def plot_progress(data, input_bounds):\n",
     "    # plot contour line and data points\n",
-    "    import matplotlib.pyplot as plt\n",
     "\n",
     "    fig = plt.figure()\n",
     "    ax = fig.add_axes([0, 0, 2, 2])\n",
@@ -363,7 +365,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 57,
+   "execution_count": null,
    "id": "04352480",
    "metadata": {},
    "outputs": [
@@ -402,6 +404,7 @@
     }
    ],
    "source": [
+    "import matplotlib.pyplot as plt\n",
     "from tqdm.notebook import tqdm\n",
     "\n",
     "for _ in tqdm(range(80)):\n",
@@ -420,7 +423,6 @@
     "\n",
     "def plot_progress(data):\n",
     "    # set up plot\n",
-    "    import matplotlib.pyplot as plt\n",
     "\n",
     "    fig = plt.figure()\n",
     "    ax = fig.add_axes([0, 0, 2, 2])\n",

docs/notebooks/trees/bo_with_trees.ipynb

@@ -15,6 +15,7 @@ dependencies = [
   "pyomo>=6.7.3",
   "onnx",
   "onnxruntime",
+  "onnxscript",
 ]
 requires-python = ">=3.9"
 dynamic = ["version"]
@@ -23,7 +24,7 @@ license = { file = "LICENSE.rst" }
 description = "OMLT is a Python package for representing machine learning models (such as neural networks) within the Pyomo optimization environment."
 
 [project.optional-dependencies]
-linear-tree = ["linear-tree"]
+linear-tree = ["linear-tree @ git+https://github.com/All-Things-Optimal/linear-tree.git@main"]
 keras = ["tensorflow", "keras"]
 keras-gpu = ["tensorflow[and-cuda]", "keras"]
 torch = ["torch", "torch-geometric"]
@@ -98,7 +99,9 @@ ignore = [
   "ANN002",
   "ANN201",
   "ANN202",
-  "ANN204"
+  "ANN204",
+  "PLW1641",
+  "TRY003"
 ]
 
 [tool.ruff.lint.pydocstyle]
@@ -171,6 +174,7 @@ module = [
   "tensorflow.*",
   "torch_geometric.*",
   "juliacall.*",
+  "onnxscript.*",
 ]
 ignore_missing_imports = true
 

pyproject.toml

@@ -30,6 +30,7 @@
 ATTR_INT = 2
 ATTR_TENSOR = 4
 ATTR_INTS = 7
+ATTR_STRING = 3
 
 
 class NetworkParser:
@@ -262,7 +263,7 @@ def _consume_dense_nodes(  # noqa: C901, PLR0912
 
         input_output_size = _get_input_output_size(input_layer, transformer)
 
-        output_size = input_output_size[:-1] + [node_weights.shape[1]]
+        output_size = [*input_output_size[:-1], node_weights.shape[1]]
 
         activation = "linear"
         if len(next_nodes) == 1:
@@ -316,12 +317,12 @@ def _consume_gemm_dense_nodes(self, node, next_nodes):
         input_output_size = _get_input_output_size(input_layer, transformer)
 
         # output is the same size as input except for the last dimension
-        output_size = input_output_size[:-1] + [weights.shape[1]]
+        output_size = [*input_output_size[:-1], weights.shape[1]]
 
         activation = "linear"
         if len(next_nodes) == 1:
             # check if Relu
-            type_, maybe_node, maybe_next_nodes = self._nodes[next_nodes[0]]
+            _, maybe_node, maybe_next_nodes = self._nodes[next_nodes[0]]
             if maybe_node.op_type in _ACTIVATION_OP_TYPES:
                 node = maybe_node
                 activation = node.op_type.lower()
@@ -376,15 +377,25 @@ def _consume_conv_nodes(self, node, next_nodes):  # noqa: PLR0912, C901, PLR0915
         biases = np.zeros(out_channels) if in_2 is None else self._initializers[in_2]
 
         attr = _collect_attributes(node)
-
+        if "strides" not in attr:
+            node_name = node.name
+            msg = f"{node_name} is missing required 'strides' attribute."
+            raise ValueError(msg)
         strides = attr["strides"]
+
         # check only kernel shape and stride are set
-        if attr["kernel_shape"] != kernel_shape:
-            msg = (
-                f"Kernel shape attribute {attr['kernel_shape']} does not match"
-                f" initialized kernel shape {kernel_shape}."
-            )
-            raise ValueError(msg)
+        if "kernel_shape" in attr:
+            if attr["kernel_shape"] != kernel_shape:
+                msg = (
+                    f"Kernel shape attribute {attr['kernel_shape']} does not match"
+                    f" initialized kernel shape {kernel_shape}."
+                )
+                raise ValueError(msg)
+        else:
+            # infer kernel shape from weights (ONNX default behavior)
+            attr["kernel_shape"] = list(kernel_shape)
+            # Assign to _kernel_shape attribute for testing purposes
+            self._kernel_shape = list(kernel_shape)
         if len(kernel_shape) != len(strides):
             msg = (
                 f"Initialized kernel shape {kernel_shape} has {len(kernel_shape)} "
@@ -437,7 +448,7 @@ def _consume_conv_nodes(self, node, next_nodes):  # noqa: PLR0912, C901, PLR0915
         activation = "linear"
         if len(next_nodes) == 1:
             # check if Relu
-            type_, maybe_node, maybe_next_nodes = self._nodes[next_nodes[0]]
+            _, maybe_node, maybe_next_nodes = self._nodes[next_nodes[0]]
             if maybe_node.op_type in _ACTIVATION_OP_TYPES:
                 node = maybe_node
                 activation = maybe_node.op_type.lower()
@@ -479,7 +490,16 @@ def _consume_reshape_nodes(self, node, next_nodes):
             raise ValueError(msg)
         [in_0, in_1] = list(node.input)
         input_layer = self._node_map[in_0]
-        new_shape = self._constants[in_1]
+        if in_1 in self._constants:
+            new_shape = self._constants[in_1]
+        elif in_1 in self._initializers:
+            new_shape = self._initializers[in_1]
+        else:
+            msg = (
+                f"Reshape node {node.name} has shape input {in_1} "
+                "that is neither a Constant nor an initializer."
+            )
+            raise KeyError(msg)
         output_size = np.empty(input_layer.output_size).reshape(new_shape).shape
         transformer = IndexMapper(input_layer.output_size, list(output_size))
         self._node_map[node.output[0]] = (transformer, input_layer)
@@ -584,7 +604,7 @@ def _consume_pool_nodes(self, node, next_nodes):  # noqa: PLR0912, C901, PLR0915
         activation = "linear"
         if len(next_nodes) == 1:
             # check if Relu
-            type_, maybe_node, maybe_next_nodes = self._nodes[next_nodes[0]]
+            _, maybe_node, maybe_next_nodes = self._nodes[next_nodes[0]]
             if maybe_node.op_type in _ACTIVATION_OP_TYPES:
                 node = maybe_node
                 activation = maybe_node.op_type.lower()
@@ -624,6 +644,8 @@ def _collect_attributes(node):
             r[attr.name] = numpy_helper.to_array(attr.t)
         elif attr.type == ATTR_INTS:  # INTS
             r[attr.name] = list(attr.ints)
+        elif attr.type == ATTR_STRING:  # STRING
+            r[attr.name] = attr.s.decode("utf-8")
         else:
             msg = f"unhandled attribute type {attr.type}"
             raise RuntimeError(msg)

src/omlt/io/onnx_parser.py

@@ -360,6 +360,9 @@ def _add_gdp_formulation_to_block(  # noqa: PLR0913 C901
     # Use the input_bounds and the linear models in the leaves to calculate
     # the lower and upper bounds on the output variable. Required for Pyomo.GDP
     scaled_output_bounds = _build_output_bounds(model_definition, scaled_input_bounds)
+    unscaled_output_bounds = _build_output_bounds(
+        model_definition, unscaled_input_bounds
+    )
 
     # Outputs are automatically scaled based on whether inputs are scaled
     for output_idx in output_indices:
@@ -374,13 +377,6 @@ def _add_gdp_formulation_to_block(  # noqa: PLR0913 C901
     max_scaled_output = np.max(scaled_output_bounds[:, 1])
     min_scaled_output = np.min(scaled_output_bounds[:, 0])
 
-    if unscaled_input_bounds is not None:
-        unscaled_output_bounds = _build_output_bounds(
-            model_definition, unscaled_input_bounds
-        )
-        block.outputs.setub(unscaled_output_bounds[1])
-        block.outputs.setlb(unscaled_output_bounds[0])
-
     if model_definition.is_scaled is True:
         block.intermediate_output = pe.Var(
             set_index, bounds=(min_scaled_output, max_scaled_output)
@@ -439,5 +435,7 @@ def disjunction_rule(b, tree):
 
     transformation_string = "gdp." + transformation
 
-    if transformation != "custom":
+    if transformation == "mbigm":
+        pe.TransformationFactory(transformation_string).apply_to(block, threads=1)
+    elif transformation != "custom":
         pe.TransformationFactory(transformation_string).apply_to(block)

src/omlt/linear_tree/lt_formulation.py

@@ -16,7 +16,7 @@ def full_space_dense_layer(net_block, net, layer_block, layer):  # noqa: PLR0912
         \end{align*}
 
     """
-    input_layer, input_layer_block = _input_layer_and_block(net_block, net, layer)
+    _, input_layer_block = _input_layer_and_block(net_block, net, layer)
 
     constraint_factory = OmltConstraintFactory()
     layer_block.dense_layer = constraint_factory.new_constraint(
@@ -105,7 +105,7 @@ def full_space_gnn_layer(net_block, net, layer_block, layer):
             \end{cases}
         \end{align*}
     """
-    input_layer, input_layer_block = _input_layer_and_block(net_block, net, layer)
+    _, input_layer_block = _input_layer_and_block(net_block, net, layer)
 
     var_factory = OmltVarFactory()
     constraint_factory = OmltConstraintFactory()
@@ -237,7 +237,7 @@ def full_space_conv2d_layer(net_block, net, layer_block, layer):
         succ_layer.activation = layer.activation
         layer.activation = "linear"
 
-    input_layer, input_layer_block = _input_layer_and_block(net_block, net, layer)
+    _, input_layer_block = _input_layer_and_block(net_block, net, layer)
     constraint_factory = OmltConstraintFactory()
 
     layer_block.convolutional_layer = constraint_factory.new_constraint(