Add LazyTensor class to implement ir.TensorProtocol

[justinchuby]

I used copilot to help implement #2231. The lazy tensor class allows users to delay transformations to the tensors until serialization time, which helps with memory usage and avoids the need to cache of unload intermediate tensor data to disk.

Example

>>> import numpy as np
>>> from onnxscript import ir
>>> weights = np.array([[1, 2, 3]])
>>> def create_tensor():
...     # Delay applying transformations to the weights
...     weights_t = weights.transpose()
...     return ir.tensor(weights_t)
>>> lazy_tensor = ir.LazyTensor(create_tensor, dtype=ir.DataType.INT64, shape=ir.Shape([1, 3]))
>>> print(lazy_tensor.numpy())
[[1]
 [2]
 [3]]
>>> print(lazy_tensor.tobytes())
b'\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00'

Fixes #2231

---

For more details, open the Copilot Workspace session.

[codecov]

❌ 3 Tests Failed:

Tests completed	Failed	Passed	Skipped
14455	3	14452	1879

View the top 3 failed test(s) by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0668_test_mod_mixed_sign_int16

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_mod_mixed_sign_int16'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_mod_mixed_sign_int16' (e=No module named 'tests.onnx_backend_test_code.test_mod_mixed_sign_int16') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_mod_mixed_sign_int16.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_mod_mixed_sign_int16.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import INT16
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_mod_mixed_sign_int16(x: INT16[6], y: INT16[6]) -> (INT16[6]):
E       z = opset13.Mod(x, y)
E       return z

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0833_test_reduce_l2_negative_axes_keep_dims_example

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_reduce_l2_negative_axes_keep_dims_example'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_reduce_l2_negative_axes_keep_dims_example' (e=No module named 'tests.onnx_backend_test_code.test_reduce_l2_negative_axes_keep_dims_example') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l2_negative_axes_keep_dims_example.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l2_negative_axes_keep_dims_example.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset18
E   
E   @script()
E   def bck_test_reduce_l2_negative_axes_keep_dims_example(data: FLOAT[3,2,2], axes: INT64[1]) -> (FLOAT[3,2,1]):
E       reduced = opset18.ReduceL2(data, axes, keepdims=1)
E       return reduced

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1005_test_scatter_elements_with_reduction_max

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_scatter_elements_with_reduction_max'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_scatter_elements_with_reduction_max' (e=No module named 'tests.onnx_backend_test_code.test_scatter_elements_with_reduction_max') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_scatter_elements_with_reduction_max.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_scatter_elements_with_reduction_max.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset18
E   
E   @script()
E   def bck_test_scatter_elements_with_reduction_max(data: FLOAT[1,5], indices: INT64[1,2], updates: FLOAT[1,2]) -> (FLOAT[1,5]):
E       y = opset18.ScatterElements(data, indices, updates, axis=1, reduction='max')
E       return y

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[Copilot]

Pull Request Overview

This PR introduces a LazyTensor class that implements the TensorProtocol and defers expensive tensor transformations until methods like numpy() or tobytes() are called. Key changes include:

• Adding the LazyTensor class and its implementation details in onnxscript/ir/_core.py.
• Implementing unit tests for LazyTensor behavior in onnxscript/ir/_core_test.py.
• Exporting LazyTensor in onnxscript/ir/init.py to make it available as part of the public API.

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
onnxscript/ir/_core_test.py	Adds tests to validate LazyTensor initialization, numpy conversion, and byte serialization.
onnxscript/ir/_core.py	Implements the LazyTensor class with support for lazy evaluation and caching.
onnxscript/ir/init.py	Exports LazyTensor for public use.

Comments suppressed due to low confidence (1)

onnxscript/ir/_core_test.py:1343

• The expected byte literal in the tobytes test is hard-coded and might be brittle on platforms with different endianness or integer representations. Consider generating the expected bytes dynamically using lazy_tensor.numpy().tobytes() to improve test resilience.

lazy_tensor.tobytes(),
            b"\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00"

[titaiwangms]

Will there be a follow up that actually uses this in any implementation?

[justinchuby]

Yes! See microsoft/onnxruntime-genai#1427