[IR] Implement register_initializer #1941

justinchuby · 2024-11-13T00:35:57Z

Implement register_initializer(value) on Graph for robustly adding an initializer to the graph. Users can also directly modify the graph.initializers dictionary, but this method does more comprehensive checks before adding, and calling this method is simpler than modifying the dictionary.

The method could be in the convenience too, but I put it here due to its relevance and for better discoverability.

codecov · 2024-11-13T00:40:30Z

❌ 16 Tests Failed:

Tests completed	Failed	Passed	Skipped
14277	16	14261	1630

View the full list of 3 ❄️ flaky tests

tests.eager_mode_test.TestEagerModeArguments_0_reference_runtime::test_function_input_and_attribute_by_kwargs_out_of_order

Flake rate in main: 38.99% (Passed 7016 times, Failed 4483 times)

Stack Traces | 0.003s run time

..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:91: in run
    res = self._run(x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:139: in _run
    res = (convert_from_ml_dtypes(res[0]),)
..../test_torch_nightly/lib/python3.12.../onnx/reference/custom_element_types.py:50: in convert_from_ml_dtypes
    return array.view(dtype=dtype)
E   ValueError: Changing the dtype of a 0d array is only supported if the itemsize is unchanged

The above exception was the direct cause of the following exception:
tests/eager_mode_test.py:115: in test_function_input_and_attribute_by_kwargs_out_of_order
    self.assertEqual(add_with_alpha(alpha=3.0, other=2.0, this=1.0), 7.0)
onnxscript/values.py:576: in __call__
    return evaluator.default().eval_function(self, args, kwargs)
onnxscript/evaluator.py:307: in eval_function
    result = function.function(*adapted_args, **adapted_kwargs)
tests/eager_mode_test.py:59: in add_with_alpha
    other = op.Mul(other, alpha)
.../onnx_opset/_impl/opset14.py:696: in Mul
    return op(*self._prepare_inputs(schema, A, B))
onnxscript/values.py:304: in __call__
    return evaluator.default().eval(schema, args, kwargs)
onnxscript/evaluator.py:194: in eval
    outputs = self._eval(schema, inputs, attributes, closure)
onnxscript/evaluator.py:524: in _eval
    result = session.run(None, session_run_input)
..../test_torch_nightly/lib/python3.12.../onnx/reference/reference_evaluator.py:599: in run
    outputs = node.run(*inputs, **linked_attributes)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:114: in run
    res = OpRunBinary.run(self, x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:93: in run
    raise TypeError(
E   TypeError: Issues with types <class 'numpy.ndarray'>, <class 'numpy.ndarray'> (binary operator 'Mul').

tests.eager_mode_test.TestEagerModeArguments_0_reference_runtime::test_function_all_input_by_kwargs

Flake rate in main: 38.99% (Passed 7016 times, Failed 4483 times)

Stack Traces | 0.003s run time

..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:91: in run
    res = self._run(x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:139: in _run
    res = (convert_from_ml_dtypes(res[0]),)
..../test_torch_nightly/lib/python3.12.../onnx/reference/custom_element_types.py:50: in convert_from_ml_dtypes
    return array.view(dtype=dtype)
E   ValueError: Changing the dtype of a 0d array is only supported if the itemsize is unchanged

The above exception was the direct cause of the following exception:
tests/eager_mode_test.py:109: in test_function_all_input_by_kwargs
    self.assertEqual(add_with_alpha(this=1.0, other=2.0), 3.0)
onnxscript/values.py:576: in __call__
    return evaluator.default().eval_function(self, args, kwargs)
onnxscript/evaluator.py:307: in eval_function
    result = function.function(*adapted_args, **adapted_kwargs)
tests/eager_mode_test.py:59: in add_with_alpha
    other = op.Mul(other, alpha)
.../onnx_opset/_impl/opset14.py:696: in Mul
    return op(*self._prepare_inputs(schema, A, B))
onnxscript/values.py:304: in __call__
    return evaluator.default().eval(schema, args, kwargs)
onnxscript/evaluator.py:194: in eval
    outputs = self._eval(schema, inputs, attributes, closure)
onnxscript/evaluator.py:524: in _eval
    result = session.run(None, session_run_input)
..../test_torch_nightly/lib/python3.12.../onnx/reference/reference_evaluator.py:599: in run
    outputs = node.run(*inputs, **linked_attributes)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:114: in run
    res = OpRunBinary.run(self, x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:93: in run
    raise TypeError(
E   TypeError: Issues with types <class 'numpy.ndarray'>, <class 'numpy.ndarray'> (binary operator 'Mul').

tests.eager_mode_test.TestEagerModeArguments_0_reference_runtime::test_function_attribute_by_positional_args

Flake rate in main: 38.99% (Passed 7016 times, Failed 4483 times)

Stack Traces | 0.003s run time

..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:91: in run
    res = self._run(x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:139: in _run
    res = (convert_from_ml_dtypes(res[0]),)
..../test_torch_nightly/lib/python3.12.../onnx/reference/custom_element_types.py:50: in convert_from_ml_dtypes
    return array.view(dtype=dtype)
E   ValueError: Changing the dtype of a 0d array is only supported if the itemsize is unchanged

The above exception was the direct cause of the following exception:
tests/eager_mode_test.py:112: in test_function_attribute_by_positional_args
    self.assertEqual(add_with_alpha(1.0, 2.0, 3.0), 7.0)
onnxscript/values.py:576: in __call__
    return evaluator.default().eval_function(self, args, kwargs)
onnxscript/evaluator.py:307: in eval_function
    result = function.function(*adapted_args, **adapted_kwargs)
tests/eager_mode_test.py:59: in add_with_alpha
    other = op.Mul(other, alpha)
.../onnx_opset/_impl/opset14.py:696: in Mul
    return op(*self._prepare_inputs(schema, A, B))
onnxscript/values.py:304: in __call__
    return evaluator.default().eval(schema, args, kwargs)
onnxscript/evaluator.py:194: in eval
    outputs = self._eval(schema, inputs, attributes, closure)
onnxscript/evaluator.py:524: in _eval
    result = session.run(None, session_run_input)
..../test_torch_nightly/lib/python3.12.../onnx/reference/reference_evaluator.py:599: in run
    outputs = node.run(*inputs, **linked_attributes)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:114: in run
    res = OpRunBinary.run(self, x, y)
..../test_torch_nightly/lib/python3.12.../reference/ops/_op.py:93: in run
    raise TypeError(
E   TypeError: Issues with types <class 'numpy.ndarray'>, <class 'numpy.ndarray'> (binary operator 'Mul').

To view more test analytics, go to the Test Analytics Dashboard
Got feedback? Let us know on Github

gramalingam · 2024-11-13T17:58:39Z

onnxscript/ir/_core.py

+            ValueError: If the value does not have its ``.const_value`` set.
+        """
+        if value.name in self._initializers:
+            if self._initializers[value.name] is not value:


Not necessarily in this PR, but would be nice to have a method to reuse initializers, especially for common duplicated constants, like 0, 1 or even shapes like (batchsize, seqlen, hiddensize).

On a related note, for very small tensors, we could replace the is comparison above with a value-equality check. May not be that important if we add a separate utility to reuse initializers as above, but if not this could serve so approximately.

I think this can be a graph pass too (dedup initializers)

gramalingam · 2024-11-13T18:01:28Z

onnxscript/ir/_core.py

+                )
+        if not value.name:
+            raise ValueError(f"Initializer must have a name: {value!r}")
+        if value.producer() is not None:


I wonder if we should support promotion of a Constant node to an initializer (later on perhaps; could be an option in this method or a separate one).

I think that can be a graph pass?

titaiwangms

Should we update this usage to exporter in terms of robustness?

justinchuby · 2024-11-13T19:21:43Z

Should we update this usage to exporter in terms of robustness?

Sure, we can do that

justinchuby added 2 commits November 13, 2024 00:29

init

f9406be

test

e1c29a0

justinchuby added topic: api topic: IR Intermediate representation labels Nov 13, 2024

justinchuby requested a review from titaiwangms November 13, 2024 00:36

justinchuby requested a review from gramalingam November 13, 2024 05:46

gramalingam reviewed Nov 13, 2024

View reviewed changes

gramalingam approved these changes Nov 13, 2024

View reviewed changes

titaiwangms approved these changes Nov 13, 2024

View reviewed changes

justinchuby enabled auto-merge (squash) November 13, 2024 19:21

justinchuby merged commit fa7d13a into main Nov 13, 2024
25 of 41 checks passed

justinchuby deleted the justinchu/register-initializer branch November 13, 2024 19:21

justinchuby mentioned this pull request Nov 13, 2024

[optimizer] Deduplicate initializer pass #1943

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[IR] Implement register_initializer #1941

[IR] Implement register_initializer #1941

justinchuby commented Nov 13, 2024 •

edited

Loading

codecov bot commented Nov 13, 2024 •

edited

Loading

gramalingam Nov 13, 2024

justinchuby Nov 13, 2024 •

edited

Loading

justinchuby Nov 13, 2024

gramalingam Nov 13, 2024

justinchuby Nov 13, 2024

justinchuby Nov 13, 2024

titaiwangms left a comment

justinchuby commented Nov 13, 2024

[IR] Implement register_initializer #1941

[IR] Implement register_initializer #1941

Conversation

justinchuby commented Nov 13, 2024 • edited Loading

codecov bot commented Nov 13, 2024 • edited Loading

❌ 16 Tests Failed:

gramalingam Nov 13, 2024

Choose a reason for hiding this comment

justinchuby Nov 13, 2024 • edited Loading

Choose a reason for hiding this comment

justinchuby Nov 13, 2024

Choose a reason for hiding this comment

gramalingam Nov 13, 2024

Choose a reason for hiding this comment

justinchuby Nov 13, 2024

Choose a reason for hiding this comment

justinchuby Nov 13, 2024

Choose a reason for hiding this comment

titaiwangms left a comment

Choose a reason for hiding this comment

justinchuby commented Nov 13, 2024

justinchuby commented Nov 13, 2024 •

edited

Loading

codecov bot commented Nov 13, 2024 •

edited

Loading

justinchuby Nov 13, 2024 •

edited

Loading