chore: import upstream snapshot with attribution

tools/Polygraphy/tests/util/test_array.py

@@ -0,0 +1,254 @@
+#
+# 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 copy
+
+import numpy as np
+import pytest
+import torch
+
+from polygraphy import cuda, util
+from polygraphy.datatype import DataType
+
+
+@pytest.mark.parametrize(
+    "obj",
+    [
+        np.transpose(np.ones((2, 3), dtype=np.float32)),
+        torch.transpose(torch.ones((2, 3), dtype=torch.float32), 1, 0),
+        cuda.DeviceArray(shape=(2, 3), dtype=DataType.FLOAT32),
+    ],
+    ids=[
+        "numpy",
+        "torch",
+        "DeviceView",
+    ],
+)
+class TestArrayFuncs:
+    def test_nbytes(self, obj):
+        nbytes = util.array.nbytes(obj)
+        assert isinstance(nbytes, int)
+        assert nbytes == 24
+
+    def test_data_ptr(self, obj):
+        data_ptr = util.array.data_ptr(obj)
+        assert isinstance(data_ptr, int)
+
+    def test_make_contiguous(self, obj):
+        if isinstance(obj, cuda.DeviceView):
+            pytest.skip("DeviceViews are always contiguous")
+
+        obj = copy.copy(obj)
+        assert not util.array.is_contiguous(obj)
+
+        obj = util.array.make_contiguous(obj)
+        assert util.array.is_contiguous(obj)
+
+    def test_dtype(self, obj):
+        assert util.array.dtype(obj) == DataType.FLOAT32
+
+    def test_view(self, obj):
+        obj = util.array.make_contiguous(obj)
+        view = util.array.view(obj, dtype=DataType.UINT8, shape=(24, 1))
+        assert util.array.dtype(view) == DataType.UINT8
+        assert util.array.shape(view) == (24, 1)
+
+    def test_resize(self, obj):
+        # Need to make a copy since we're modifying the array.
+        obj = copy.copy(util.array.make_contiguous(obj))
+        obj = util.array.resize_or_reallocate(obj, (1, 1))
+        assert util.array.shape(obj) == (1, 1)
+
+
+@pytest.mark.parametrize(
+    "obj, is_on_cpu",
+    [
+        (np.ones((2, 3)), True),
+        (torch.ones((2, 3)), True),
+        (torch.ones((2, 3), device="cuda"), False),
+        (cuda.DeviceArray(shape=(2, 3), dtype=DataType.FLOAT32), False),
+    ],
+)
+def test_is_on_cpu(obj, is_on_cpu):
+    assert util.array.is_on_cpu(obj) == is_on_cpu
+
+
+@pytest.mark.parametrize(
+    "obj, is_on_gpu",
+    [
+        (np.ones((2, 3)), False),
+        (torch.ones((2, 3)), False),
+        (torch.ones((2, 3), device="cuda"), True),
+        (cuda.DeviceArray(shape=(2, 3), dtype=DataType.FLOAT32), True),
+    ],
+)
+def test_is_on_cpu(obj, is_on_gpu):
+    assert util.array.is_on_gpu(obj) == is_on_gpu
+
+
+@pytest.mark.parametrize(
+    "lhs,rhs,expected",
+    [
+        (np.ones((2, 3)), np.ones((2, 3)), True),
+        (np.zeros((2, 3)), np.ones((2, 3)), False),
+        (torch.ones((2, 3)), torch.ones((2, 3)), True),
+        (torch.zeros((2, 3)), torch.ones((2, 3)), False),
+    ],
+)
+def test_equal(lhs, rhs, expected):
+    assert util.array.equal(lhs, rhs) == expected
+
+
+@pytest.mark.parametrize(
+    "index,shape",
+    [
+        (7, (4, 4)),
+        (12, (4, 4, 3, 2)),
+    ],
+)
+def test_unravel_index(index, shape):
+    assert util.array.unravel_index(index, shape) == np.unravel_index(index, shape)
+
+
+@pytest.mark.parametrize(
+    "lhs, rhs, expected",
+    [
+        (np.array([5.00001]), np.array([5.00]), True),
+        (np.array([5.5]), np.array([5.00]), False),
+        (torch.tensor([5.00001]), torch.tensor([5.00]), True),
+        (torch.tensor([5.5]), torch.tensor([5.00]), False),
+    ],
+)
+def test_allclose(lhs, rhs, expected):
+    assert util.array.allclose(lhs, rhs) == expected
+
+
+ARRAYS = [
+    # Generate ints so FP rounding error is less of an issue
+    np.random.randint(1, 25, size=(5, 2)).astype(np.float32),
+    # Make sure functions work with an even or odd number of elements
+    np.random.randint(1, 25, size=(1, 3)).astype(np.float32),
+    # Generate binary values
+    np.random.randint(0, 2, size=(5, 2)).astype(np.float32),
+    # Test with scalars
+    np.ones(shape=tuple(), dtype=np.float32),
+]
+
+TEST_CASES = []
+IDS = []
+for arr in ARRAYS:
+    TEST_CASES.extend([(arr, arr), (torch.from_numpy(arr), arr)])
+    IDS.extend(["numpy", "torch"])
+
+
+@pytest.mark.parametrize("obj, np_arr", TEST_CASES, ids=IDS)
+class TestArrayMathFuncs:
+    # Test that the util.array implementations match NumPy
+    @pytest.mark.parametrize(
+        "func, np_func",
+        [
+            (util.array.max, np.amax),
+            (util.array.argmax, np.argmax),
+            (util.array.min, np.amin),
+            (util.array.argmin, np.argmin),
+            (util.array.mean, np.mean),
+            (util.array.std, np.std),
+            (util.array.var, np.var),
+            (util.array.median, np.median),
+            (util.array.any, np.any),
+            (util.array.all, np.all),
+        ],
+    )
+    def test_reduction_funcs(self, obj, np_arr, func, np_func):
+        assert np.isclose(func(obj), np_func(np_arr))
+
+    @pytest.mark.parametrize(
+        "func, np_func",
+        [
+            (util.array.abs, np.abs),
+            (util.array.isinf, np.isinf),
+            (util.array.isnan, np.isnan),
+            (util.array.argwhere, np.argwhere),
+        ],
+    )
+    def test_array_funcs(self, obj, np_arr, func, np_func):
+        obj = func(obj)
+        assert util.array.equal(obj, np.array(np_func(np_arr)))
+
+    def test_cast(self, obj, np_arr):
+        dtype = DataType.INT32
+        casted = util.array.cast(obj, dtype)
+        assert util.array.dtype(casted) == dtype
+        assert type(casted) == type(obj)
+
+    def test_to_torch(self, obj, np_arr):
+        assert isinstance(util.array.to_torch(obj), torch.Tensor)
+
+    def test_to_numpy(self, obj, np_arr):
+        assert isinstance(util.array.to_numpy(obj), np.ndarray)
+
+    def test_histogram(self, obj, np_arr):
+        hist, bins = util.array.histogram(obj)
+        np_hist, np_bins = np.histogram(np_arr)
+        np_hist = np_hist.astype(np_arr.dtype)
+
+        assert util.array.allclose(hist, np_hist)
+        assert util.array.allclose(bins, np_bins)
+
+    @pytest.mark.parametrize("k", [1, 2, 3, 4])
+    @pytest.mark.parametrize("axis", [0, 1])
+    def test_topk(self, obj, np_arr, k, axis):
+        if axis >= len(util.array.shape(obj)):
+            pytest.skip()
+        topk_vals = util.array.topk(obj, k, axis)
+
+        k_clamped = min(util.array.shape(obj)[axis], k)
+        tensor = util.array.to_torch(np_arr)
+        ref_topk_vals = torch.topk(tensor, k_clamped, axis)
+
+        assert util.array.allclose(topk_vals[0], ref_topk_vals[0])
+
+    @pytest.mark.parametrize(
+        "func, np_func",
+        [
+            (util.array.subtract, np.subtract),
+            (util.array.divide, np.divide),
+            (util.array.logical_xor, np.logical_xor),
+            (util.array.logical_and, np.logical_and),
+            (util.array.greater, np.greater),
+        ],
+    )
+    def test_binary_funcs(self, obj, np_arr, func, np_func):
+        obj = func(obj, obj + 1)
+        assert util.array.equal(obj, np.array(np_func(np_arr, np_arr + 1)))
+
+    @pytest.mark.parametrize(
+        "func, np_func, types",
+        [
+            (
+                util.array.where,
+                np.where,
+                tuple(map(DataType.from_dtype, (np.bool8, np.float32, np.float32))),
+            ),
+        ],
+    )
+    def test_ternary_funcs(self, obj, np_arr, func, np_func, types):
+        build_inputs = lambda input: map(
+            lambda pair: util.array.cast(input + pair[0], pair[1]), enumerate(types)
+        )
+        obj = func(*build_inputs(obj))
+        assert util.array.equal(obj, np.array(np_func(*build_inputs(np_arr))))

tools/Polygraphy/tests/util/test_serde.py

@@ -0,0 +1,218 @@
+#
+# 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 numpy as np
+import pytest
+import tensorrt as trt
+import torch
+
+from polygraphy import constants, util
+from polygraphy.backend.trt import Algorithm, TacticReplayData, TensorInfo
+from polygraphy.comparator import IterationResult, RunResults
+from polygraphy.exception import PolygraphyException
+from polygraphy.json import Decoder, Encoder, from_json, load_json, to_json
+
+
+class Dummy:
+    def __init__(self, x):
+        self.x = x
+
+
+@Encoder.register(Dummy)
+def encode_dummy(dummy):
+    return {"x": dummy.x}
+
+
+@Decoder.register(Dummy)
+def decode_dummy(dct):
+    assert len(dct) == 1  # Custom type markers should be removed at this point
+    return Dummy(x=dct["x"])
+
+
+class NoDecoder:
+    def __init__(self, x):
+        self.x = x
+
+
+@Encoder.register(NoDecoder)
+def encode_nodecoder(no_decoder):
+    return {"x": no_decoder.x}
+
+
+class TestEncoder:
+    def test_registered(self):
+        d = Dummy(x=-1)
+        d_json = to_json(d)
+        assert encode_dummy(d) == {"x": d.x, constants.TYPE_MARKER: "Dummy"}
+        expected = f'{{\n    "x": {d.x},\n    "{constants.TYPE_MARKER}": "Dummy"\n}}'
+        assert d_json == expected
+
+
+class TestDecoder:
+    def test_object_pairs_hook(self):
+        d = Dummy(x=-1)
+        d_json = to_json(d)
+
+        new_d = from_json(d_json)
+        assert new_d.x == d.x
+
+    def test_error_on_no_decoder(self):
+        d = NoDecoder(x=1)
+        d_json = to_json(d)
+
+        with pytest.raises(
+            PolygraphyException,
+            match="Could not decode serialized type: NoDecoder. This could be because a required module is missing.",
+        ):
+            from_json(d_json)
+
+    def test_names_correct(self):
+        # Trigger `try_register_common_json`
+        d = Dummy(x=-1)
+        to_json(d)
+
+        # If the name of a class changes, then we need to specify an `alias` when registering
+        # to retain backwards compatibility.
+        assert set(Decoder.polygraphy_registered.keys()) == {
+            "__polygraphy_encoded_Algorithm",
+            "__polygraphy_encoded_AttentionLayerHint",
+            "__polygraphy_encoded_Dummy",
+            "__polygraphy_encoded_FormattedArray",
+            "__polygraphy_encoded_IterationContext",
+            "__polygraphy_encoded_IterationResult",
+            "__polygraphy_encoded_LazyArray",
+            "__polygraphy_encoded_ndarray",
+            "__polygraphy_encoded_RunResults",
+            "__polygraphy_encoded_ShardHints",
+            "__polygraphy_encoded_ShardTensor",
+            "__polygraphy_encoded_TacticReplayData",
+            "__polygraphy_encoded_Tensor",
+            "__polygraphy_encoded_TensorInfo",
+            "Algorithm",
+            "AttentionLayerHint",
+            "Dummy",
+            "FormattedArray",
+            "IterationContext",
+            "IterationResult",
+            "LazyArray",
+            "LazyNumpyArray",
+            "ndarray",
+            "RunResults",
+            "ShardHints",
+            "ShardTensor",
+            "TacticReplayData",
+            "Tensor",
+            "TensorInfo",
+        }
+
+
+def make_algo():
+    return Algorithm(
+        implementation=4,
+        tactic=5,
+        # Should work even if strides are not set
+        inputs=[
+            TensorInfo(trt.float32, (1, 2), -1, 1),
+            TensorInfo(trt.float32, (1, 2), -1, 1),
+        ],
+        outputs=[TensorInfo(trt.float32, (2, 3), -1, 1)],
+    )
+
+
+def make_iter_result():
+    return IterationResult(
+        runtime=4.5,
+        runner_name="test",
+        outputs={
+            "out0": np.random.random_sample((1, 2, 1)),
+            "out1": np.ones((1, 2), dtype=np.float32),
+        },
+    )
+
+
+JSONABLE_CASES = [
+    RunResults([("runner0", [make_iter_result()]), ("runner0", [make_iter_result()])]),
+    TacticReplayData().add("hi", algorithm=make_algo()),
+]
+
+
+class TestImplementations:
+    @pytest.mark.parametrize(
+        "obj",
+        [
+            TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+            Algorithm(
+                implementation=4,
+                tactic=5,
+                inputs=[TensorInfo(trt.float32, (1, 2, 3), -1, 1)],
+                outputs=[TensorInfo(trt.float32, (1, 2, 3), -1, 1)],
+            ),
+            Algorithm(
+                implementation=4,
+                tactic=5,
+                inputs=[
+                    TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                    TensorInfo(trt.int8, (1, 2, 3), -1, 1),
+                ],
+                outputs=[TensorInfo(trt.float16, (1, 2, 3), -1, 1)],
+            ),
+            np.ones((3, 4, 5), dtype=np.int64),
+            np.ones(5, dtype=np.int64),
+            np.zeros((4, 5), dtype=np.float32),
+            np.random.random_sample((3, 5)),
+            torch.ones((3, 4, 5), dtype=torch.int64),
+            make_iter_result(),
+            RunResults(
+                [("runner0", [make_iter_result()]), ("runner0", [make_iter_result()])]
+            ),
+        ],
+        ids=lambda x: type(x),
+    )
+    def test_serde(self, obj):
+        encoded = to_json(obj)
+        decoded = from_json(encoded)
+        if isinstance(obj, np.ndarray):
+            assert np.array_equal(decoded, obj)
+        elif isinstance(obj, torch.Tensor):
+            assert torch.equal(decoded, obj)
+        else:
+            assert decoded == obj
+
+    @pytest.mark.parametrize("obj", JSONABLE_CASES)
+    def test_to_from_json(self, obj):
+        encoded = obj.to_json()
+        decoded = type(obj).from_json(encoded)
+        assert decoded == obj
+
+    @pytest.mark.parametrize("obj", JSONABLE_CASES)
+    def test_save_load(self, obj):
+        with util.NamedTemporaryFile("w+") as f:
+            obj.save(f)
+            decoded = type(obj).load(f)
+            assert decoded == obj
+
+    def test_cannot_save_load_to_different_types(self):
+        run_result = JSONABLE_CASES[0]
+        encoded = run_result.to_json()
+
+        with pytest.raises(PolygraphyException, match="JSON cannot be decoded into"):
+            TacticReplayData.from_json(encoded)
+
+
+def test_load_json_errors_if_file_nonexistent():
+    with pytest.raises(FileNotFoundError, match="No such file"):
+        load_json("polygraphy-nonexistent-path")

tools/Polygraphy/tests/util/test_util.py

@@ -0,0 +1,420 @@
+#
+# 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 contextlib
+import io
+import os
+import random
+import tempfile
+from multiprocessing import Process
+
+import numpy as np
+import pytest
+
+from polygraphy import util
+from polygraphy.backend.trt import engine_from_network, network_from_onnx_bytes
+from polygraphy.util import util as util_internal  # For accessing and testing private functions in util.py
+from tests.models.meta import ONNX_MODELS
+
+VOLUME_CASES = [
+    ((1, 1, 1), 1),
+    ((2, 3, 4), 24),
+    (tuple(), 1),
+]
+
+
+@pytest.mark.parametrize("case", VOLUME_CASES)
+def test_volume(case):
+    it, vol = case
+    assert util.volume(it) == vol
+
+
+class FindStrInIterableCase:
+    def __init__(self, name, seq, index, expected):
+        self.name = name
+        self.seq = seq
+        self.index = index
+        self.expected = expected
+
+
+FIND_STR_IN_ITERABLE_CASES = [
+    # Case insensitve, plus function should return element from sequence, not name.
+    FindStrInIterableCase(
+        "Softmax:0", seq=["Softmax:0"], index=None, expected="Softmax:0"
+    ),
+    FindStrInIterableCase(
+        "Softmax:0", seq=["softmax:0"], index=None, expected="softmax:0"
+    ),
+    # Exact matches should take priority
+    FindStrInIterableCase(
+        "exact_name",
+        seq=["exact_name_plus", "exact_name"],
+        index=0,
+        expected="exact_name",
+    ),
+    # Index should come into play when no matches are found
+    FindStrInIterableCase(
+        "non-existent", seq=["test", "test2"], index=1, expected="test2"
+    ),
+]
+
+
+@pytest.mark.parametrize("case", FIND_STR_IN_ITERABLE_CASES)
+def test_find_str_in_iterable(case):
+    actual = util.find_str_in_iterable(case.name, case.seq, case.index)
+    assert actual == case.expected
+
+
+SHAPE_OVERRIDE_CASES = [
+    ((1, 3, 224, 224), (None, 3, 224, 224), True),
+]
+
+
+@pytest.mark.parametrize("case", SHAPE_OVERRIDE_CASES)
+def test_is_valid_shape_override(case):
+    override, shape, expected = case
+    assert (
+        util.is_valid_shape_override(new_shape=override, original_shape=shape)
+        == expected
+    )
+
+
+def arange(shape):
+    return np.arange(util.volume(shape)).reshape(shape)
+
+
+SHAPE_MATCHING_CASES = [
+    (arange((1, 1, 3, 3)), (3, 3), arange((3, 3))),  # Squeeze array shape
+    (
+        arange((1, 3, 3, 1)),
+        (1, 1, 3, 3),
+        arange((1, 1, 3, 3)),
+    ),  # Permutation should make no difference as other dimensions are 1s
+    (arange((3, 3)), (1, 1, 3, 3), arange((1, 1, 3, 3))),  # Unsqueeze where needed
+    (arange((3, 3)), (-1, 3), arange((3, 3))),  # Infer dynamic
+    (
+        arange((3 * 2 * 2,)),
+        (None, 3, 2, 2),
+        arange((1, 3, 2, 2)),
+    ),  # Reshape with inferred dimension
+    (
+        arange((1, 3, 2, 2)),
+        (None, 2, 2, 3),
+        np.transpose(arange((1, 3, 2, 2)), [0, 2, 3, 1]),
+    ),  # Permute
+]
+
+build_torch = lambda a, **kwargs: util.array.to_torch(np.array(a, **kwargs))
+
+
+@pytest.mark.parametrize("array_type", [np.array, build_torch])
+@pytest.mark.parametrize("arr, shape, expected", SHAPE_MATCHING_CASES)
+def test_shape_matching(arr, shape, expected, array_type):
+    arr = util.try_match_shape(array_type(arr), shape)
+    assert util.array.equal(arr, array_type(expected))
+
+
+UNPACK_ARGS_CASES = [
+    ((0, 1, 2), 3, (0, 1, 2)),  # no extras
+    ((0, 1, 2), 4, (0, 1, 2, None)),  # 1 extra
+    ((0, 1, 2), 2, (0, 1)),  # 1 fewer
+]
+
+
+@pytest.mark.parametrize("case", UNPACK_ARGS_CASES)
+def test_unpack_args(case):
+    args, num, expected = case
+    assert util.unpack_args(args, num) == expected
+
+
+UNIQUE_LIST_CASES = [
+    ([], []),
+    ([3, 1, 2], [3, 1, 2]),
+    ([1, 2, 3, 2, 1], [1, 2, 3]),
+    ([0, 0, 0, 0, 1, 0, 0], [0, 1]),
+    ([5, 5, 5, 5, 5], [5]),
+]
+
+
+@pytest.mark.parametrize("case", UNIQUE_LIST_CASES)
+def test_unique_list(case):
+    lst, expected = case
+    assert util.unique_list(lst) == expected
+
+
+def test_find_in_dirs():
+    with tempfile.TemporaryDirectory() as topdir:
+        dirs = list(
+            map(
+                lambda x: os.path.join(topdir, x),
+                ["test0", "test1", "test2", "test3", "test4"],
+            )
+        )
+        for subdir in dirs:
+            os.makedirs(subdir)
+
+        path_dir = random.choice(dirs)
+        path = os.path.join(path_dir, "cudart64_11.dll")
+
+        with open(path, "w") as f:
+            f.write("This file should be found by find_in_dirs")
+
+        assert util.find_in_dirs("cudart64_*.dll", dirs) == [path]
+
+
+@pytest.mark.parametrize(
+    "val,key,default,expected",
+    [
+        (1.0, None, None, 1.0),  # Basic
+        ({"inp": "hi"}, "inp", "", "hi"),  # Per-key
+        ({"inp": "hi"}, "out", "default", "default"),  # Per-key missing
+        ({"inp": 1.0, "": 2.0}, "out", 1.5, 2.0),  # Per-key with default
+    ],
+)
+def test_value_or_from_dict(val, key, default, expected):
+    actual = util.value_or_from_dict(val, key, default)
+    assert actual == expected
+
+
+def test_atomic_open():
+    def write_to_file(path, content):
+        with util.LockFile(path):
+            old_contents = util.load_file(path, mode="r")
+            util.save_file(old_contents + content, path, mode="w")
+
+    NUM_LINES = 10
+    NUM_PROCESSES = 5
+
+    outfile = util.NamedTemporaryFile()
+
+    processes = [
+        Process(
+            target=write_to_file,
+            args=(outfile.name, f"{proc} - writing line\n" * NUM_LINES),
+        )
+        for proc in range(NUM_PROCESSES)
+    ]
+
+    for process in processes:
+        process.start()
+
+    for process in processes:
+        process.join()
+
+    for process in processes:
+        assert not process.is_alive()
+        assert process.exitcode == 0
+
+    # Since we write atomically, all processes should be able to write their
+    # contents. Furthermore, the contents should be grouped by process.
+    with open(outfile.name) as f:
+        lines = list(f.readlines())
+        assert len(lines) == NUM_LINES * NUM_PROCESSES
+
+        for idx in range(NUM_PROCESSES):
+            offset = idx * NUM_LINES
+            expected_prefix = lines[offset].partition("-")[0].strip()
+            assert all(
+                line.startswith(expected_prefix)
+                for line in lines[offset : offset + NUM_LINES]
+            )
+
+    # Make sure the lock file is written to the correct path and not removed automatically.
+    assert os.path.exists(outfile.name + ".lock")
+
+
+class TestMakeRepr:
+    def test_basic(self):
+        assert util.make_repr("Example", 1, x=2) == ("Example(1, x=2)", False, False)
+
+    def test_default_args(self):
+        assert util.make_repr("Example", None, None, x=2) == (
+            "Example(None, None, x=2)",
+            True,
+            False,
+        )
+
+    def test_empty_args_are_default(self):
+        assert util.make_repr("Example", x=2) == ("Example(x=2)", True, False)
+
+    def test_default_kwargs(self):
+        assert util.make_repr("Example", 1, 2, x=None, y=None) == (
+            "Example(1, 2)",
+            False,
+            True,
+        )
+
+    def test_empty_kwargs_are_default(self):
+        assert util.make_repr("Example", 1, 2) == ("Example(1, 2)", False, True)
+
+    def test_does_not_modify(self):
+        obj = {"x": float("inf")}
+        assert util.make_repr("Example", obj) == (
+            "Example({'x': float('inf')})",
+            False,
+            True,
+        )
+        assert obj == {"x": float("inf")}
+
+    @pytest.mark.parametrize("obj", [float("nan"), float("inf"), float("-inf")])
+    @pytest.mark.parametrize("recursion_depth", [0, 1, 2])
+    def test_nan_inf(self, obj, recursion_depth):
+        if obj == float("inf"):
+            expected = "float('inf')"
+        elif obj == float("-inf"):
+            expected = "float('-inf')"
+        else:
+            expected = "float('nan')"
+
+        for _ in range(recursion_depth):
+            obj = {"x": obj}
+            expected = f"{{'x': {expected}}}"
+
+        assert util.make_repr("Example", obj) == (f"Example({expected})", False, True)
+
+
+@pytest.mark.serial
+def test_check_called_by():
+    outfile = io.StringIO()
+    with contextlib.redirect_stdout(outfile):
+        warn_msg = "Calling 'test_check_called_by.<locals>.callee()' directly is not recommended. Please use 'caller()' instead."
+
+        @util.check_called_by("caller")
+        def callee():
+            pass
+
+        def caller():
+            return callee()
+
+        # If we call via the caller, no message should be emitted
+        caller()
+        outfile.seek(0)
+        out = outfile.read()
+        assert warn_msg not in out
+
+        # If we call the callee directly, we should see a warning
+        callee()
+        outfile.seek(0)
+        out = outfile.read()
+        assert warn_msg in out
+
+
+class TestGetNumBytes:
+    def test_should_get_given_str(self) -> None:
+        """Test that _get_num_bytes returns the correct number of bytes when given `str`."""
+        # Precondition.
+        contents = "hello"
+
+        # Under test.
+        num_bytes = util_internal._get_num_bytes(contents)
+
+        # Postcondition.
+        assert num_bytes == len("hello")
+
+    def test_should_get_given_bytes(self) -> None:
+        """Test that _get_num_bytes returns the correct number of bytes when given `bytes`."""
+        # Precondition.
+        contents = bytes(b"hello")
+
+        # Under test.
+        num_bytes = util_internal._get_num_bytes(contents)
+
+        # Postcondition.
+        assert num_bytes == len("hello")
+
+    def test_should_get_given_IHostMemory(self) -> None:
+        """Test that _get_num_bytes returns the correct number of bytes when given `IHostMemory`."""
+        # Precondition.
+        contents = engine_from_network(network_from_onnx_bytes(ONNX_MODELS["identity"].loader)).serialize()
+
+        # Under test.
+        num_bytes = util_internal._get_num_bytes(contents)
+
+        # Postcondition.
+        assert num_bytes == len(memoryview(contents))
+
+    def test_should_raise_error_given_invalid_type(self) -> None:
+        """Test that _get_num_bytes raises an error when given an invalid type."""
+        # Precondition.
+        invalid_contents = 123
+
+        # Under test and postcondition.
+        with pytest.raises(
+            TypeError, match=f"`contents` is {invalid_contents}, which is not bytes-like. Cannot get number of bytes."
+        ):
+            util_internal._get_num_bytes(invalid_contents)
+
+
+class TestSaveFile:
+    def test_should_save_str_to_path(self) -> None:
+        """Test that `save_file` should save a string to a path."""
+        # Precondition.
+        contents = "hello"
+        with util.NamedTemporaryFile("w+") as f:
+            dest = f.name
+
+            # Under test.
+            util.save_file(contents, dest, mode="w+")
+
+            # Postcondition.
+            f.seek(0)
+            written_contents = f.read()
+            assert contents == written_contents
+
+    def test_should_save_str_to_file_like(self) -> None:
+        """Test that `save_file` should save a string to a file-like object."""
+        # Precondition.
+        contents = "hello"
+        with util.NamedTemporaryFile("w+") as f:
+            dest = f
+
+            # Under test.
+            util.save_file(contents, dest, mode="w+")
+
+            # Postcondition.
+            f.seek(0)
+            written_contents = f.read()
+            assert contents == written_contents
+
+    def test_should_save_bytes_to_path(self) -> None:
+        """Test that `save_file` should save bytes to a path."""
+        # Precondition.
+        contents = b"hello"
+        with util.NamedTemporaryFile("wb+") as f:
+            dest = f.name
+
+            # Under test.
+            util.save_file(contents, dest, mode="wb+")
+
+            # Postcondition.
+            f.seek(0)
+            written_contents = f.read()
+            assert contents == written_contents
+
+    def test_should_save_IHostMemory_to_path(self) -> None:
+        """Test that `save_file` should save an `IHostMemory` to a path."""
+        # Precondition.
+        contents = engine_from_network(network_from_onnx_bytes(ONNX_MODELS["identity"].loader)).serialize()
+        with util.NamedTemporaryFile("wb+") as f:
+            dest = f.name
+
+            # Under test.
+            util.save_file(contents, dest, mode="wb+")
+
+            # Postcondition.
+            f.seek(0)
+            written_contents = f.read()
+            assert bytes(contents) == written_contents