#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import tempfile

import onnx
from onnx import numpy_helper
import onnx_graphsurgeon as gs
import pytest
from polygraphy import util
from polygraphy.backend.onnx import util as onnx_util
from polygraphy.tools.sparse import SparsityPruner
from tests.helper import is_file_non_empty
from tests.models.meta import ONNX_MODELS, model_path


@pytest.fixture()
def onnx_model_sanity_check(poly_run):
    def onnx_model_sanity_check_impl(model_path):
        poly_run([model_path, "--model-type=onnx", "--onnxrt"])

    return onnx_model_sanity_check_impl


def get_exclude_list(exclude_list):
    if not exclude_list:
        return set()
    with open(exclude_list) as fp:
        lines = [line.rstrip() for line in fp]
        return set(lines)


def pruned_initializer_sanity_check(opath, is_sparse=False, exclude_list=None):
    exclude_list = get_exclude_list(exclude_list)
    # we only prune the input data of QuantizeLinear and leave the scale and zero_point untouched
    if "qdq" in opath:
        exclude_list.add("y_scale")
        exclude_list.add("y_zero_point")

    model = onnx.load(opath)
    for initializer in model.graph.initializer:
        # If initializer is to be left un-stripped
        if initializer.name in exclude_list:
            # ensure initializer is non-empty and doc_string doesn't contain the weightless flag
            shape_match = (
                list(numpy_helper.to_array(initializer).shape) == initializer.dims
            )
            if "TRT_WEIGHTLESS" in initializer.doc_string or not shape_match:
                return False
            continue

        # ensure initializer is empty and doc_string is in required format
        init_empty = initializer.raw_data == b""
        trt_weightless, sparsity = initializer.doc_string.split("/")
        trt_weightless_correctness = trt_weightless == "TRT_WEIGHTLESS"
        sparsity_correctness = False
        if (not is_sparse and sparsity == "") or (
            is_sparse and sparsity == "SPARSE_2_4"
        ):
            sparsity_correctness = True

        if not (init_empty and trt_weightless_correctness and sparsity_correctness):
            return False

    return True


def get_initializers_to_sparsify(ipath):
    model = onnx.load(ipath)
    initializers_to_sparsify = set()
    for initializer in model.graph.initializer:
        if "SPARSE_2_4" in initializer.doc_string:
            initializers_to_sparsify.add(initializer.name)

    return initializers_to_sparsify


def reconstructed_initializer_sanity_check(opath, initializers_to_sparsify):
    model = onnx.load(opath)
    sparsity_checker = SparsityPruner(model)
    sparsity_checker.check()
    sparse_tensors = sparsity_checker.sparse_tensors
    for initializer in model.graph.initializer:
        shape_match = list(numpy_helper.to_array(initializer).shape) == initializer.dims
        if not shape_match:
            return False

        # ensure sparsity of initializers is retained
        if (
            initializer.name in initializers_to_sparsify
            and initializer.name not in sparse_tensors
        ):
            return False

    return True


def was_shape_inference_run(status, model):
    logging_correct = "Shape inference completed successfully" in (
        status.stdout + status.stderr
    )

    has_shape = True
    model = onnx.load(model)
    for val in model.graph.value_info:
        has_shape &= val.type.tensor_type.HasField("shape")

    return logging_correct and has_shape


class TestSurgeonExtract:
    def test_no_shape_inference_if_has_metadata(
        self, poly_surgeon_extract, onnx_model_sanity_check
    ):
        with util.NamedTemporaryFile() as outmodel:
            status = poly_surgeon_extract(
                [
                    ONNX_MODELS["identity_identity"].path,
                    "-o",
                    outmodel.name,
                    "--inputs",
                    "X:auto:auto",
                ]
            )
            onnx_model_sanity_check(outmodel.name)
            assert not was_shape_inference_run(status, outmodel.name)

    def test_onnx_shape_inference_if_no_metadata(
        self, poly_surgeon_extract, onnx_model_sanity_check
    ):
        with util.NamedTemporaryFile() as outmodel:
            status = poly_surgeon_extract(
                [
                    ONNX_MODELS["identity_identity"].path,
                    "-o",
                    outmodel.name,
                    "--inputs",
                    "identity_out_0:auto:auto",
                ]
            )
            onnx_model_sanity_check(outmodel.name)
            assert was_shape_inference_run(status, outmodel.name)

    def test_fallback_shape_inference_no_onnx_shape_inference(
        self, poly_surgeon_extract, onnx_model_sanity_check
    ):
        with util.NamedTemporaryFile() as outmodel:
            status = poly_surgeon_extract(
                [
                    ONNX_MODELS["identity_identity"].path,
                    "-o",
                    outmodel.name,
                    "--inputs",
                    "identity_out_0:auto:auto",
                    "--outputs",
                    "identity_out_2:auto",
                    "--force-fallback-shape-inference",
                ]
            )
            onnx_model_sanity_check(outmodel.name)
            assert not was_shape_inference_run(status, outmodel.name)

    def test_force_fallback_shape_inference_will_override_model_shapes(
        self, poly_surgeon_extract, onnx_model_sanity_check
    ):
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon_extract(
                [
                    ONNX_MODELS["dynamic_identity"].path,
                    "-o",
                    outmodel.name,
                    "--outputs",
                    "Y:auto",
                    "--force-fallback-shape-inference",
                ]
            )
            onnx_model_sanity_check(outmodel.name)
            graph = gs.import_onnx(onnx.load(outmodel.name))
            # Inputs should become fixed since fallback shape inference is being forced.
            for tensor in graph.tensors().values():
                assert tensor.shape is not None
            assert tuple(graph.inputs[0].shape) == (1, 2, 1, 1)
            assert tuple(graph.outputs[0].shape) == (1, 2, 1, 1)

    def test_sanity_dim_param(self, poly_surgeon_extract, onnx_model_sanity_check):
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon_extract([ONNX_MODELS["dim_param"].path, "-o", outmodel.name])
            onnx_model_sanity_check(outmodel.name)


class TestSurgeonInsert:
    def check_insert_model(
        self,
        path,
        expected_node_ops,
        expected_graph_input_names,
        expected_graph_output_names,
    ):
        model = onnx.load(path)
        assert [node.op_type for node in model.graph.node] == expected_node_ops

        graph_output_names = set([out.name for out in model.graph.output])
        assert graph_output_names == set(expected_graph_output_names)
        assert all(onnx_util.get_dtype(out) for out in model.graph.output)

        graph_input_names = set([out.name for out in model.graph.input])
        assert graph_input_names == set(expected_graph_input_names)
        return model

    def test_insert_at_tensor(self, poly_surgeon):
        # Insert a new node in between existing nodes without replacing any existing nodes.
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon(
                [
                    "insert",
                    ONNX_MODELS["identity_identity"].path,
                    "-o",
                    outmodel.name,
                    "--inputs=identity_out_0",
                    "--outputs=identity_out_0",
                    "--op=FakeOp",
                ]
            )
            self.check_insert_model(
                outmodel.name,
                ["Identity", "FakeOp", "Identity"],
                ["X"],
                ["identity_out_2"],
            )

    def test_graph_output(self, poly_surgeon):
        # FakeOp output tensor should be marked as a graph output. Name should be preserved - identity_out_2
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon(
                [
                    "insert",
                    ONNX_MODELS["identity_identity"].path,
                    "-o",
                    outmodel.name,
                    "--inputs=identity_out_2",
                    "--outputs=identity_out_2",
                    "--op=FakeOp",
                ]
            )
            self.check_insert_model(
                outmodel.name,
                ["Identity", "Identity", "FakeOp"],
                ["X"],
                ["identity_out_2"],
            )

    def test_at_graph_input(self, poly_surgeon):
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon(
                [
                    "insert",
                    ONNX_MODELS["identity_identity"].path,
                    "-o",
                    outmodel.name,
                    "--inputs=X",
                    "--outputs=X",
                    "--op=FakeOp",
                ]
            )
            self.check_insert_model(
                outmodel.name,
                ["FakeOp", "Identity", "Identity"],
                ["X"],
                ["identity_out_2"],
            )

    # When a specified input tensor is used by multiple other nodes, it should not be
    # disconnected from other nodes.
    def test_multi_use_input_replace_single_path(self, poly_surgeon):
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon(
                [
                    "insert",
                    ONNX_MODELS["reducable"].path,
                    "-o",
                    outmodel.name,
                    "--inputs=add_out_4",
                    "--outputs=identity_out_8",
                    "--op=FakeOp",
                ]
            )
            model = self.check_insert_model(
                outmodel.name,
                ["Identity", "Identity", "Add", "FakeOp", "Identity"],
                ["X0", "Y0"],
                ["identity_out_6", "identity_out_8"],
            )
            other_branch_node = model.graph.node[-1]
            assert other_branch_node.name == "onnx_graphsurgeon_node_7"
            assert other_branch_node.input == ["add_out_4"]

    def test_multi_use_input_replace_all_paths(self, poly_surgeon):
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon(
                [
                    "insert",
                    ONNX_MODELS["reducable"].path,
                    "-o",
                    outmodel.name,
                    "--inputs=add_out_4",
                    "--outputs=add_out_4",
                    "--op=FakeOp",
                ]
            )
            model = self.check_insert_model(
                outmodel.name,
                ["Identity", "Identity", "Add", "FakeOp", "Identity", "Identity"],
                ["X0", "Y0"],
                ["identity_out_6", "identity_out_8"],
            )
            for branch_num in range(1, 3):
                branch = model.graph.node[-branch_num]
                assert branch.input == ["add_out_4_polygraphy_surgeon_insert_output"]

    def test_with_attributes(self, poly_surgeon):
        with util.NamedTemporaryFile() as outmodel:
            # str_attr='0' should be interpreted as a string, not an int
            # float_attr=0.0 should be interpreted as a float, not an int
            # int_attr=0 should be interpreted as an int
            poly_surgeon(
                [
                    "insert",
                    ONNX_MODELS["identity_identity"].path,
                    "-o",
                    outmodel.name,
                    "--inputs=X",
                    "--outputs=X",
                    "--op=FakeOp",
                    "--attrs",
                    "str_attr='0'",
                    "int_attr=0",
                    "float_attr=0.0",
                    "other_str_attr=name",
                    "str_list_attr=['0','1']",
                    "int_list_attr=[1,2,3]",
                    "float_list_attr=[0.0,-1.0,-2.0]",
                ]
            )
            model = self.check_insert_model(
                outmodel.name,
                ["FakeOp", "Identity", "Identity"],
                ["X"],
                ["identity_out_2"],
            )
            node = model.graph.node[0]

            attrs = node.attribute
            assert attrs[0].name == "str_attr"
            assert attrs[0].s == b"0"

            assert attrs[1].name == "int_attr"
            assert attrs[1].i == 0

            assert attrs[2].name == "float_attr"
            assert attrs[2].f == 0.0

            assert attrs[3].name == "other_str_attr"
            assert attrs[3].s == b"name"

            assert attrs[4].name == "str_list_attr"
            assert attrs[4].strings == [b"0", b"1"]

            assert attrs[5].name == "int_list_attr"
            assert attrs[5].ints == [1, 2, 3]

            assert attrs[6].name == "float_list_attr"
            assert attrs[6].floats == [0.0, -1.0, -2.0]


class TestSurgeonSanitize:
    @pytest.mark.parametrize(
        "no_per_pass_shape_inf", [None, "--no-per-pass-shape-inference"]
    )
    @pytest.mark.parametrize("fold_shapes", [None, "--no-fold-shapes"])
    @pytest.mark.parametrize("partitioning", [None, "basic", "recursive"])
    @pytest.mark.parametrize(
        "no_onnxruntime_shape_inference", [None, "--no-onnxruntime-shape-inference"]
    )
    def test_fold_constants(
        self,
        poly_surgeon,
        no_per_pass_shape_inf,
        partitioning,
        fold_shapes,
        onnx_model_sanity_check,
        no_onnxruntime_shape_inference,
    ):
        with util.NamedTemporaryFile() as outmodel:
            cmd = [
                "sanitize",
                ONNX_MODELS["const_foldable"].path,
                "-o",
                outmodel.name,
                "--fold-constants",
                "-v",
            ]
            if fold_shapes:
                cmd += [fold_shapes]
            if partitioning:
                cmd += ["--partitioning", partitioning]
            if no_per_pass_shape_inf:
                cmd += [no_per_pass_shape_inf]
            if no_onnxruntime_shape_inference:
                cmd += [no_onnxruntime_shape_inference]
            status = poly_surgeon(cmd)

            assert (
                "Inferred shapes in the model with `onnxruntime.tools.symbolic_shape_infer`"
                in status.stdout
            ) == (no_onnxruntime_shape_inference is None)

            onnx_model_sanity_check(outmodel.name)
            model = onnx.load(outmodel.name)
            assert len(model.graph.node) == 1

    @pytest.mark.parametrize("global_upper_bound", [None, "2000"])
    @pytest.mark.parametrize("specified_upper_bound", [None, "cast_out_6:4000"])
    def test_set_upper_bound(
        self,
        poly_surgeon,
        global_upper_bound,
        specified_upper_bound,
        onnx_model_sanity_check,
    ):
        with util.NamedTemporaryFile() as outmodel:
            cmd = [
                "sanitize",
                ONNX_MODELS["unbounded_dds"].path,
                "-o",
                outmodel.name,
                "--set-unbounded-dds-upper-bound",
            ]
            upper_bound = "1000"
            if global_upper_bound:
                upper_bound = "2000"
                cmd += [global_upper_bound]
            if specified_upper_bound:
                upper_bound = "4000"
                cmd += [specified_upper_bound]
            if global_upper_bound is None and specified_upper_bound is None:
                cmd += [upper_bound]
            poly_surgeon(cmd)

            onnx_model_sanity_check(outmodel.name)
            graph = gs.import_onnx(onnx.load(outmodel.name))

            # Check if there is a Min operator in the modified model
            find_min = False
            for node in graph.nodes:
                if node.op == "Min":
                    find_min = True
                    # Check if the Min operator's second input is a constant tensor
                    assert isinstance(node.inputs[1], gs.Constant)

                    val = node.inputs[1].values
                    # Check if the constant value equals the target upper bound
                    assert str(val) == upper_bound
            assert find_min

    def test_fold_constants_single_pass(self, poly_surgeon, onnx_model_sanity_check):
        with util.NamedTemporaryFile() as outmodel:
            status = poly_surgeon(
                [
                    "sanitize",
                    ONNX_MODELS["const_foldable"].path,
                    "-o",
                    outmodel.name,
                    "--fold-constants",
                    "--num-passes=1",
                ]
            )

            assert "Pass 1" in status.stdout
            assert "Pass 2" not in status.stdout

            onnx_model_sanity_check(outmodel.name)
            model = onnx.load(outmodel.name)
            assert len(model.graph.node) == 1

    @pytest.mark.parametrize("new_dim", [1, 2, 3])
    def test_override_shapes(self, poly_surgeon, new_dim, onnx_model_sanity_check):
        with util.NamedTemporaryFile() as outmodel:
            cmd = [
                "sanitize",
                ONNX_MODELS["dynamic_identity"].path,
                "-o",
                outmodel.name,
                f"--override-input-shapes=X:[1,2,{new_dim},{new_dim}]",
            ]
            poly_surgeon(cmd)
            onnx_model_sanity_check(outmodel.name)

            model = onnx.load(outmodel.name)

            shape = []
            for dim in model.graph.input[0].type.tensor_type.shape.dim:
                assert isinstance(dim.dim_value, int) and dim.dim_value >= 0
                shape.append(dim.dim_value)

            assert shape == [1, 2, new_dim, new_dim]

    def test_override_shapes_no_clear_const_tensors_meta(self, poly_surgeon):
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon(
                [
                    "sanitize",
                    ONNX_MODELS["const_foldable"].path,
                    "-o",
                    outmodel.name,
                    "--override-input-shapes=input:[1,3]",
                ]
            )

    def test_override_shapes_partial_inputs(self, poly_surgeon):
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon(
                [
                    "sanitize",
                    ONNX_MODELS["dynamic_identity"].path,
                    "-o",
                    outmodel.name,
                    "--override-input-shapes=Y:[1,2,3,4]",
                ]
            )
            model = onnx.load(outmodel.name)
            assert (
                model.graph.input[0].type.tensor_type.shape.dim[2].dim_param == "height"
            )
            assert (
                model.graph.input[0].type.tensor_type.shape.dim[3].dim_param == "width"
            )

    def test_override_shapes_no_reorder(self, poly_surgeon):
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon(
                [
                    "sanitize",
                    ONNX_MODELS["reducable"].path,
                    "-o",
                    outmodel.name,
                    "--override-input-shapes",
                    "Y0:[5]",
                    "X0:[5]",
                ]
            )
            model = onnx.load(outmodel.name)
            assert model.graph.input[0].name == "X0"
            assert model.graph.input[1].name == "Y0"

    def test_modify_onnx_outputs(self, poly_surgeon):
        with util.NamedTemporaryFile(suffix=".onnx") as outmodel:
            poly_surgeon(
                [
                    "sanitize",
                    ONNX_MODELS["identity_identity"].path,
                    "-o",
                    outmodel.name,
                    "--outputs",
                    "mark",
                    "all",
                ]
            )

            model = onnx.load(outmodel.name)
            assert len(model.graph.output) == 2

    def test_cleanup(self, poly_surgeon):
        with util.NamedTemporaryFile(suffix=".onnx") as outmodel:
            poly_surgeon(
                [
                    "sanitize",
                    ONNX_MODELS["identity_identity"].path,
                    "-o",
                    outmodel.name,
                    "--outputs",
                    "identity_out_0",
                    "--cleanup",
                ]
            )

            model = onnx.load(outmodel.name)
            assert len(model.graph.node) == 1
            assert model.graph.output[0].name == "identity_out_0"

    def test_toposort(self, poly_surgeon):
        with util.NamedTemporaryFile(suffix=".onnx") as outmodel:
            poly_surgeon(
                [
                    "sanitize",
                    ONNX_MODELS["unsorted"].path,
                    "-o",
                    outmodel.name,
                    "--toposort",
                ]
            )

            model = onnx.load(outmodel.name)
            assert model.graph.node[0].name == "onnx_graphsurgeon_node_1"
            assert model.graph.node[1].name == "onnx_graphsurgeon_node_3"

    def test_external_data(self, poly_surgeon, poly_run):
        with tempfile.TemporaryDirectory() as outdir:
            model = ONNX_MODELS["ext_weights"]
            outmodel = os.path.join(outdir, "out_model.onnx")
            outdata = "ext_weights.data"
            assert poly_surgeon(
                [
                    "sanitize",
                    model.path,
                    "--external-data-dir",
                    model.ext_data,
                    "--fold-constants",
                    "-o",
                    outmodel,
                    "--save-external-data",
                    outdata,
                    "--external-data-size-threshold=0",
                    "-vvvvv",
                ]
            )
            assert is_file_non_empty(outmodel)
            assert is_file_non_empty(os.path.join(outdir, outdata))
            assert poly_run([outmodel, "--onnxrt", "--external-data-dir", outdir])

    def test_force_fallback_shape_inference_will_override_model_shapes(
        self, poly_surgeon, onnx_model_sanity_check
    ):
        with util.NamedTemporaryFile() as outmodel:
            status = poly_surgeon(
                [
                    "sanitize",
                    ONNX_MODELS["dynamic_identity"].path,
                    "-o",
                    outmodel.name,
                    "--force-fallback-shape-inference",
                ]
            )
            onnx_model_sanity_check(outmodel.name)
            graph = gs.import_onnx(onnx.load(outmodel.name))
            # Inputs should become fixed since fallback shape inference is being forced.
            assert tuple(graph.inputs[0].shape) == (1, 2, 1, 1)
            for tensor in graph.tensors().values():
                assert tensor.shape is not None
            assert tuple(graph.outputs[0].shape) == (1, 2, 1, 1)

            assert not was_shape_inference_run(status, outmodel.name)

    @pytest.mark.parametrize(
        "size_threshold, expect_folding",
        [
            ("10M", True),
            ("9.99M", False),
        ],
    )
    def test_size_threshold(
        self, poly_surgeon, size_threshold, expect_folding, onnx_model_sanity_check
    ):
        with util.NamedTemporaryFile() as outmodel:
            poly_surgeon(
                [
                    "sanitize",
                    ONNX_MODELS["constant_fold_bloater"].path,
                    "-o",
                    outmodel.name,
                    "--fold-constants",
                    "--fold-size-threshold",
                    size_threshold,
                ]
            )

            onnx_model_sanity_check(outmodel.name)
            model = onnx.load(outmodel)

            if expect_folding:
                assert len(model.graph.node) == 0
            else:
                assert len(model.graph.node) == 1
                assert model.graph.node[0].op_type == "Tile"


class TestSurgeonPrune:
    @pytest.mark.parametrize(
        "model_name",
        ["matmul", "matmul.fp16", "matmul.bf16", "matmul.bf16.i32data", "conv"],
    )
    def test_prune(self, poly_surgeon, onnx_model_sanity_check, model_name):
        with tempfile.TemporaryDirectory() as outdir:
            ipath = ONNX_MODELS[model_name].path
            opath = os.path.join(outdir, "pruned." + os.path.basename(ipath))
            status = poly_surgeon(["prune", ipath, "-o", opath])
            assert status
            if "bf16" not in ipath:
                onnx_model_sanity_check(opath)


class TestSurgeonWeightStrip:
    @pytest.mark.parametrize(
        "model_name",
        [
            "matmul",
            "matmul.fp16",
            "matmul.bf16",
            "conv",
            "sparse.matmul",
            "sparse.conv",
            "transpose_matmul",
            "qdq_conv",
        ],
    )
    def test_weight_strip(self, poly_surgeon, model_name):
        with tempfile.TemporaryDirectory() as outdir:
            ipath = ONNX_MODELS[model_name].path
            opath = os.path.join(outdir, "weightless." + os.path.basename(ipath))
            status = poly_surgeon(["weight-strip", ipath, "-o", opath])
            assert status

            is_sparse = "sparse" in ipath
            assert pruned_initializer_sanity_check(opath, is_sparse=is_sparse)

    @pytest.mark.parametrize(
        "model_name,     exclude_list",
        [
            ["matmul", "matmul.exclude_list.txt"],
            ["sparse.conv", "sparse.conv.exclude_list.txt"],
            ["qdq_conv", "qdq_conv.exclude_list.txt"],
        ],
    )
    def test_weight_strip_exclude_file(self, poly_surgeon, model_name, exclude_list):
        with tempfile.TemporaryDirectory() as outdir:
            ipath = ONNX_MODELS[model_name].path
            exclude_list = model_path(exclude_list)
            opath = os.path.join(outdir, "weightless_sparse." + os.path.basename(ipath))
            status = poly_surgeon(
                ["weight-strip", ipath, "-o", opath, "--exclude-list", exclude_list]
            )
            assert status

            is_sparse = "sparse" in ipath
            assert pruned_initializer_sanity_check(
                opath, is_sparse=is_sparse, exclude_list=exclude_list
            )


class TestSurgeonWeightReconstruct:
    @pytest.mark.parametrize(
        "model_name",
        [
            "weightless.matmul.fp16",
            "weightless.matmul.bf16",
            "weightless.conv",
            "weightless.sparse.matmul",
            "weightless.sparse.conv",
            "weightless.transpose_matmul",
            "weightless.qdq_conv",
        ],
    )
    def test_weight_reconstruct(
        self, poly_surgeon, onnx_model_sanity_check, model_name
    ):
        with tempfile.TemporaryDirectory() as outdir:
            ipath = ONNX_MODELS[model_name].path
            opath = os.path.join(outdir, "reconstruct." + os.path.basename(ipath))
            status = poly_surgeon(["weight-reconstruct", ipath, "-o", opath])
            assert status
            if "bf16" not in ipath:
                onnx_model_sanity_check(opath)

                initializers_to_sparsify = get_initializers_to_sparsify(ipath)
                assert reconstructed_initializer_sanity_check(
                    opath, initializers_to_sparsify
                )
