From df82a822684c6d3f7e22b7f386b55e17cbda4439 Mon Sep 17 00:00:00 2001 From: Colin Peppler Date: Thu, 12 Oct 2023 10:54:55 -0700 Subject: [PATCH 1/2] Temporary commit at 10/5/2023, 2:10:31 PM Differential Revision: D49759512 test_codegen fbshipit-source-id: c2d9f9ee4aeda516e56229b5369bf46cc313e256 --- python/aitemplate/backend/codegen.py | 46 +-- .../backend/test_codegen_output_aliases.py | 267 ++++++++++++++++++ .../backend/test_codegen_output_tensor.py | 174 ++++++++++++ 3 files changed, 468 insertions(+), 19 deletions(-) create mode 100644 tests/unittest/backend/test_codegen_output_aliases.py create mode 100644 tests/unittest/backend/test_codegen_output_tensor.py diff --git a/python/aitemplate/backend/codegen.py b/python/aitemplate/backend/codegen.py index 4ec52c311..e3b79cafe 100644 --- a/python/aitemplate/backend/codegen.py +++ b/python/aitemplate/backend/codegen.py @@ -370,7 +370,6 @@ def __init__( self.visited_func = set() self.visited_dims = set() self.set_up_constant_names = [] - self.param_name_to_ptr_idx = {} self.num_constants = 0 self.constants_data_size = 0 @@ -584,7 +583,6 @@ def _codegen_input_tensor(self, tensor: Tensor) -> None: ) ) self.set_inputs.append(check_not_null(tensor)) - self.param_name_to_ptr_idx[name] = self.input_idx self._record_param_tensor_info(tensor, self.input_idx) self.input_idx += 1 @@ -599,29 +597,27 @@ def _get_output_idx(self, name: str) -> int: def _codegen_output_aliases_tensor(self, tensor: Tensor) -> None: name = tensor._attrs["name"] view = tensor._attrs["is_view_of"] - if tensor._attrs["external_tensor"] is not None: + external_tensor = tensor._attrs["external_tensor"] + if external_tensor is not None: + assert not external_tensor._attrs[ + "is_param" + ], "Views of constants should be folded." self.set_inputs.append(set_value(name, view._attrs["name"])) return - is_view = view is not None - if is_view and (view._attrs["name"] in self.param_name_to_ptr_idx): - ptr_idx = self.param_name_to_ptr_idx[view._attrs["name"]] + + if view: + # View is already initialized, assign to view. self.set_inputs.append(set_value(name, view._attrs["name"])) else: - ptr_idx = self._get_output_idx(name) + # Original tensor, initialize it. + output_idx = self._get_output_idx(name) self.set_inputs.append( set_value( name, - f"static_cast(params_[{ptr_idx}].ptr)", + f"static_cast(params_[{output_idx}].ptr)", ) ) - self.param_name_to_ptr_idx[name] = ptr_idx - if tensor._attrs["is_output"]: - self._record_param_tensor_info(tensor, ptr_idx) - self.set_inputs.append( - check_not_null(tensor, skip_if_lower_bound_is_zero=True) - ) - def _codegen_output_tensor(self, tensor: Tensor) -> None: is_param = tensor._attrs["is_param"] is_input = tensor._attrs["is_input"] @@ -634,7 +630,6 @@ def _codegen_output_tensor(self, tensor: Tensor) -> None: if is_param: self._codegen_param_setup(tensor) - self._record_param_tensor_info(tensor, output_idx) self.device_to_device_copies.append(device_copy(tensor, tensor, output_idx)) elif external_tensor is not None: # Special view cases for outputs; we can hit this case if the output @@ -649,7 +644,6 @@ def _codegen_output_tensor(self, tensor: Tensor) -> None: self.device_to_device_copies.append( device_copy(tensor, external_tensor, output_idx) ) - self._record_param_tensor_info(tensor, output_idx) elif is_input: # Inputs that are also outputs require an extra copy self.set_inputs.append( @@ -659,11 +653,25 @@ def _codegen_output_tensor(self, tensor: Tensor) -> None: ) ) self._record_param_tensor_info(tensor, self.input_idx) - self._record_param_tensor_info(tensor, output_idx) self.device_to_device_copies.append(device_copy(tensor, tensor, output_idx)) self.input_idx += 1 + elif is_view: + self.set_inputs.append(set_value(name, view._attrs["name"])) + self.set_inputs.append( + check_not_null(tensor, skip_if_lower_bound_is_zero=True) + ) else: - self._codegen_output_aliases_tensor(tensor) + self.set_inputs.append( + set_value( + name, + f"static_cast(params_[{output_idx}].ptr)", + ) + ) + self.set_inputs.append( + check_not_null(tensor, skip_if_lower_bound_is_zero=True) + ) + + self._record_param_tensor_info(tensor, output_idx) def _process_dims(self, shape: List[IntVar]) -> None: for dim in shape: diff --git a/tests/unittest/backend/test_codegen_output_aliases.py b/tests/unittest/backend/test_codegen_output_aliases.py new file mode 100644 index 000000000..42e099dc5 --- /dev/null +++ b/tests/unittest/backend/test_codegen_output_aliases.py @@ -0,0 +1,267 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# 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 unittest +from typing import Sequence + +from aitemplate.backend.codegen import set_value + +from aitemplate.compiler import compile_model, ops +from aitemplate.compiler.base import Tensor +from aitemplate.compiler.ops.common.epilogue import FuncEnum +from aitemplate.testing import detect_target + +from aitemplate.testing.test_utils import gen_input_tensor, get_random_torch_tensor + + +class TestCodegenOutputAliases(unittest.TestCase): + MODEL_GENERATED_FILE = "model-generated.h" + WORKDIR = "/tmp/ait/test_codegen_output_aliases" + SHAPE = (2, 3, 4) + + def _check_string_in_file(self, file_path: str, target_string: str) -> bool: + with open(file_path, "r") as f: + content = f.read() + return target_string in content + + def _assert_codegen_exists( + self, + test_name: str, + expected_codegen: Sequence[str], + filename: str, + exists: bool = True, + ) -> bool: + model_path = os.path.join(self.WORKDIR, test_name, filename) + for line in expected_codegen: + self.assertEqual( + exists, + self._check_string_in_file(model_path, line), + f"Expected this to be codegen'd in {model_path}: \n {line}", + ) + + """ + ### Test Cases: + NOTE: <-- represents a "is_view_of" edge + 1. gelu <-- output + 2. intermediate0 <-- intermediate1 <-- output + 3. gelu <-- output_0 + ^---- output_1 + 4. input <-- intermediate <-- output + 5. constant <-- intermediate <-- output + 6. output_0 <-- intermediate <-- output_1 + """ + + def test_simple_output_alias(self, test_name="simple_output_alias"): + """ + Case: Simple case where an output aliases an internal tensor. + Graph: ( gelu ) <--view-- ( output ) + """ + x = gen_input_tensor(shape=self.SHAPE, name="input_x") + gelu = ops.elementwise(FuncEnum.GELU)(x) # has an output alias + output0 = ops.flatten()(gelu) + output0._attrs["is_output"] = True + output0._attrs["name"] = "output_0" + + compile_model( + output0, + detect_target(), + self.WORKDIR, + test_name, + do_optimize_graph=False, + ) + + view_name = output0._attrs["is_view_of"]._attrs["name"] + expected_codegen = ( + set_value( + view_name, + f"static_cast(params_[1].ptr)", + ), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_GENERATED_FILE + ) + + def test_view_of_view(self, test_name="view_of_view"): + """ + Case: There's two intermediate views, where one is a view of the other. + Graph: ( input ) <--view-- ( intermediate0 ) <--view-- ( intermediate1 ) <--view-- ( output_0 ) + """ + x = gen_input_tensor(shape=self.SHAPE, name="input_x") + intermediate0 = ops.unsqueeze(dim=0)(x) # has an output alias + intermediate1 = ops.unsqueeze(dim=0)(intermediate0) # has an output alias + output0 = ops.flatten()(intermediate1) + output0._attrs["is_output"] = True + output0._attrs["name"] = "output_0" + + compile_model( + output0, + detect_target(), + self.WORKDIR, + test_name, + do_optimize_graph=False, + ) + + intermediate0_name = intermediate0._attrs["name"] + intermediate1_name = intermediate1._attrs["name"] + expected_codegen = ( + set_value(intermediate0_name, "input_x"), + set_value(intermediate1_name, intermediate0_name), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_GENERATED_FILE + ) + + def test_double_alias(self, test_name="double_alias"): + """ + Case: Two outputs are a view of the same tensor. + Graph: ( gelu ) <--view-- ( output_0 ) + <--view-- ( output_1 ) + Expected: gelu will be initialized with params_[2] + """ + # AIT, two outputs. + x = gen_input_tensor(shape=self.SHAPE, name="input_x") + gelu = ops.elementwise(FuncEnum.GELU)(x) # has an output alias + output0 = ops.unsqueeze(dim=0)(gelu) + output0._attrs["is_output"] = True + output0._attrs["name"] = "output_0" + + output1 = ops.flatten()(gelu) + output1._attrs["is_output"] = True + output1._attrs["name"] = "output_1" + + compile_model( + [output0, output1], + detect_target(), + self.WORKDIR, + test_name, + do_optimize_graph=True, + ) + + view_name = output0._attrs["is_view_of"]._attrs["name"] + expected_codegen = ( + # input_x assigned to params_[0].ptr + set_value( + view_name, + f"static_cast(params_[2].ptr)", + # see _construct_output_name_to_index_map for params indexing + ), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_GENERATED_FILE + ) + + # ======================================================================== # + # Cases where the has_output_alias tensor also has an external tensor # + # ======================================================================== # + def test_external_tensor_is_input(self, test_name="external_tensor_is_input"): + """ + Case: The intermediate tensor is a view of an input. + Graph: ( input ) <--view-- ( intermediate ) <--view-- ( output_0 ) + """ + + x = gen_input_tensor(shape=self.SHAPE, name="input_x") + intermediate_view = ops.unsqueeze(dim=0)(x) # has an output alias + output0 = ops.flatten()(intermediate_view) + output0._attrs["is_output"] = True + output0._attrs["name"] = "output_0" + + compile_model( + output0, + detect_target(), + self.WORKDIR, + test_name, + do_optimize_graph=False, + ) + intermediate_view_name = intermediate_view._attrs["name"] + expected_codegen = ( + # input_x is assigned to params_[0].ptr + set_value(intermediate_view_name, "input_x"), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_GENERATED_FILE + ) + + def test_external_tensor_is_constant(self, test_name="external_tensor_is_constant"): + """ + Case: The intermediate tensor is a view of a constant. + Graph: ( constant ) <--view-- ( intermediate ) <--view-- ( output_0 ) + Expect: The intermediate tensor is constant folded. + """ + constant = Tensor( + shape=self.SHAPE, dtype="float16", name="constant" + ) # has an output alias + intermediate_view = ops.unsqueeze(dim=0)(constant) # has an output alias + output0 = ops.flatten()(intermediate_view) + output0._attrs["is_output"] = True + output0._attrs["name"] = "output_0" + + constant_pt = get_random_torch_tensor(self.SHAPE, dtype="float16") + + compile_model( + output0, + detect_target(), + self.WORKDIR, + test_name, + constants={"constant": constant_pt}, + do_optimize_graph=False, + ) + intermediate_view_name = intermediate_view._attrs["name"] + expected_codegen = ( + set_value(intermediate_view_name, "constant"), + set_value("output_0", intermediate_view_name), + ) + # We expect the intermediate tensor and output_0 to be constant folded. + # So, we don't expect _codegen_output_aliases_tensor to be called. + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_GENERATED_FILE, exists=False + ) + + def test_external_tensor_is_output(self, test_name="external_tensor_is_output"): + """ + Case: The intermediate tensor is a view of an output. + Graph: ( output_0 ) <--view-- ( intermediate ) <--view-- ( output_1 ) + """ + x = gen_input_tensor(shape=self.SHAPE, name="input_x") + z = gen_input_tensor(shape=self.SHAPE, name="input_z") + output0 = ops.elementwise(FuncEnum.RELU)(x + z) + output0._attrs["is_output"] = True + output0._attrs["name"] = "output_0" + + intermediate_view = ops.unsqueeze(dim=0)(output0) # has an output alias + output1 = ops.flatten()(intermediate_view) + output1._attrs["is_output"] = True + output1._attrs["name"] = "output_1" + + compile_model( + [output0, output1], + detect_target(), + self.WORKDIR, + test_name, + do_optimize_graph=False, + ) + intermediate_name = intermediate_view._attrs["name"] + expected_codegen = ( + # input_z is assigned to params_[0].ptr + # input_x is assigned to params_[1].ptr + set_value( + "output_0", + "static_cast(params_[2].ptr)", + ), + set_value(intermediate_name, "output_0"), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_GENERATED_FILE + ) diff --git a/tests/unittest/backend/test_codegen_output_tensor.py b/tests/unittest/backend/test_codegen_output_tensor.py new file mode 100644 index 000000000..cefe3ab74 --- /dev/null +++ b/tests/unittest/backend/test_codegen_output_tensor.py @@ -0,0 +1,174 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# 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 unittest +from typing import Sequence + +from aitemplate.backend.codegen import device_copy, set_value + +from aitemplate.compiler import compile_model, ops +from aitemplate.compiler.ops.common.epilogue import FuncEnum +from aitemplate.testing import detect_target + +from aitemplate.testing.test_utils import gen_input_tensor + + +class TestCodegenOutput(unittest.TestCase): + MODEL_CONTAINER_FILE = "model_container_base.cu" + MODEL_GENERATED_FILE = "model-generated.h" + WORKDIR = "/tmp/ait/test_codegen_output" + SHAPE = (2, 3, 4) + + def _check_string_in_file(self, file_path: str, target_string: str) -> bool: + with open(file_path, "r") as f: + content = f.read() + return target_string in content + + def _assert_codegen_exists( + self, + test_name: str, + expected_codegen: Sequence[str], + filename: str, + exists: bool = True, + ) -> bool: + model_path = os.path.join(self.WORKDIR, test_name, filename) + for line in expected_codegen: + self.assertEqual( + exists, + self._check_string_in_file(model_path, line), + f"Expected this to be codegen'd in {model_path}: \n {line}", + ) + + """ + ### Test Cases: + 1. gelu <-- output + 2. gelu <--view-- output_0 + ^----view-- output_1 + 3. gelu <--view-- output_0 <--view-- output_1 + """ + + def test_no_view(self, test_name="no_view"): + """ + Case: Simple case with one output. + """ + x = gen_input_tensor(shape=self.SHAPE, name="input_x") + gelu = ops.elementwise(FuncEnum.GELU)(x) + output = ops.flatten()(gelu) + output._attrs["is_output"] = True + output._attrs["name"] = "output_0" + + compile_model( + [output], + detect_target(), + self.WORKDIR, + test_name, + do_optimize_graph=True, + ) + + expected_codegen = ( + # input_x is assigned to params_[0].ptr + set_value("param_names_[1]", '"output_0"'), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_CONTAINER_FILE + ) + + expected_codegen = ( + set_value( + "output_0", + "static_cast(params_[1].ptr)", + ), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_GENERATED_FILE + ) + + def test_double_alias(self, test_name="double_alias"): + """ + Case: Two outputs are a view of the same tensor. + Graph: ( gelu ) <--view-- ( output_0 ) + <--view-- ( output_1 ) + """ + # AIT, two outputs. + x = gen_input_tensor(shape=self.SHAPE, name="input_x") + gelu = ops.elementwise(FuncEnum.GELU)(x) # has an output alias + output0 = ops.unsqueeze(dim=0)(gelu) + output0._attrs["is_output"] = True + output0._attrs["name"] = "output_0" + + output1 = ops.flatten()(gelu) + output1._attrs["is_output"] = True + output1._attrs["name"] = "output_1" + + compile_model( + [output0, output1], + detect_target(), + self.WORKDIR, + test_name, + do_optimize_graph=False, + ) + + expected_codegen = ( + set_value("param_names_[1]", '"output_0"'), + set_value("param_names_[2]", '"output_1"'), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_CONTAINER_FILE + ) + + view = output0._attrs["is_view_of"] + view_name = view._attrs["name"] + expected_codegen = ( + set_value("output_0", view_name), + set_value("output_1", view_name), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_GENERATED_FILE + ) + + def test_output_is_view_of_output(self, test_name="output_is_view_of_output"): + """ + Case: An output is a view of an output. + Graph: ( gelu ) <--view-- ( output_0 ) <--view-- ( output_1 ) + """ + x = gen_input_tensor(shape=self.SHAPE, name="input_x") + gelu = ops.elementwise(FuncEnum.GELU)(x) # has an output alias + output0 = ops.flatten()(gelu) + output0._attrs["is_output"] = True + output0._attrs["name"] = "output_0" + + output1 = ops.flatten()(output0) + output1._attrs["is_output"] = True + output1._attrs["name"] = "output_1" + + compile_model( + [output0, output1], + detect_target(), + self.WORKDIR, + test_name, + do_optimize_graph=False, + ) + + view_name = output0._attrs["is_view_of"]._attrs["name"] + expected_codegen = ( + set_value("output_0", view_name), + set_value("output_1", "output_0"), + device_copy(output1, output0, dst_idx=2), + ) + self._assert_codegen_exists( + test_name, expected_codegen, self.MODEL_GENERATED_FILE + ) From 2b87bc16c64bf79638cc807a6b91f1210b2993ae Mon Sep 17 00:00:00 2001 From: Colin Peppler Date: Thu, 12 Oct 2023 10:55:14 -0700 Subject: [PATCH 2/2] Fix accuracy issue for double output alias (#950) Summary: Pull Request resolved: https://github.com/facebookincubator/AITemplate/pull/950 ## Problem Here's an edge case for AIT. Suppose we have two outputs, and both are view on the same tensor. Atm, AIT will not provide accurate results for output0. ``` some-tensor <--view-- output0 ^------view-- ouptut1 void SetUpInputOutput() { input_x = static_cast(params_[0].ptr); elementwise_0_0 = static_cast(params_[2].ptr); output_0 = elementwise_0_0; output_1 = elementwise_0_0; } void DeviceToDeviceCopies(stream) { // empty } ``` Why doesn't AIT provide accurate results for output0? Because notice how `params_[1]` isn't assigned to anything. ## Solution Use a D2D copy to pass data from `params_[2]` to `params_[1]`. We do this by checking to see if the view is aliased by another output. * If yes, then run a D2D copy. * If no, don't worry about this output. ## Refactor We refactor `_codegen_output_tensor` by combining the `external_tensor` case with the `is_view` case. Differential Revision: D50202241 fbshipit-source-id: 599cc609dcb11f92e432479f2f3214236f6d68c5 --- python/aitemplate/backend/codegen.py | 27 ++++++++++--------- .../backend/test_codegen_output_tensor.py | 23 ++++++++++++++-- 2 files changed, 35 insertions(+), 15 deletions(-) diff --git a/python/aitemplate/backend/codegen.py b/python/aitemplate/backend/codegen.py index e3b79cafe..8353fe93f 100644 --- a/python/aitemplate/backend/codegen.py +++ b/python/aitemplate/backend/codegen.py @@ -631,19 +631,25 @@ def _codegen_output_tensor(self, tensor: Tensor) -> None: if is_param: self._codegen_param_setup(tensor) self.device_to_device_copies.append(device_copy(tensor, tensor, output_idx)) - elif external_tensor is not None: - # Special view cases for outputs; we can hit this case if the output - # is a view of a constant, input, or another output. + elif is_view or external_tensor is not None: assert ( is_view - ), f"orig_tensor is not None, but node {name} is not marked as a view! Node: {tensor}" + ), f"External tensor is not None, but node {name} is not marked as a view! Node: {tensor}" + view_name = view._attrs["name"] + self.set_inputs.append(set_value(name, view_name)) self.set_inputs.append( - check_not_null(tensor, output_idx, skip_if_lower_bound_is_zero=True) + check_not_null(tensor, skip_if_lower_bound_is_zero=True) ) - self.set_inputs.append(set_value(name, view._attrs["name"])) - self.device_to_device_copies.append( - device_copy(tensor, external_tensor, output_idx) + + view_assigned_to_another_output = ( + self._get_output_idx(view_name) != output_idx ) + if external_tensor or view_assigned_to_another_output: + # Copy from original tensor so this output can also have the data. + original_tensor = external_tensor if external_tensor else view + self.device_to_device_copies.append( + device_copy(tensor, original_tensor, output_idx) + ) elif is_input: # Inputs that are also outputs require an extra copy self.set_inputs.append( @@ -655,11 +661,6 @@ def _codegen_output_tensor(self, tensor: Tensor) -> None: self._record_param_tensor_info(tensor, self.input_idx) self.device_to_device_copies.append(device_copy(tensor, tensor, output_idx)) self.input_idx += 1 - elif is_view: - self.set_inputs.append(set_value(name, view._attrs["name"])) - self.set_inputs.append( - check_not_null(tensor, skip_if_lower_bound_is_zero=True) - ) else: self.set_inputs.append( set_value( diff --git a/tests/unittest/backend/test_codegen_output_tensor.py b/tests/unittest/backend/test_codegen_output_tensor.py index cefe3ab74..1ec21ec26 100644 --- a/tests/unittest/backend/test_codegen_output_tensor.py +++ b/tests/unittest/backend/test_codegen_output_tensor.py @@ -18,13 +18,15 @@ import unittest from typing import Sequence +import torch + from aitemplate.backend.codegen import device_copy, set_value from aitemplate.compiler import compile_model, ops from aitemplate.compiler.ops.common.epilogue import FuncEnum from aitemplate.testing import detect_target -from aitemplate.testing.test_utils import gen_input_tensor +from aitemplate.testing.test_utils import gen_input_tensor, get_random_torch_tensor class TestCodegenOutput(unittest.TestCase): @@ -102,6 +104,8 @@ def test_double_alias(self, test_name="double_alias"): Case: Two outputs are a view of the same tensor. Graph: ( gelu ) <--view-- ( output_0 ) <--view-- ( output_1 ) + Expect: If a tensor is a view for multiple outputs, then it's assigned to + only one of the outputs' ptrs. We expect D2D copies for the remaining outputs. """ # AIT, two outputs. x = gen_input_tensor(shape=self.SHAPE, name="input_x") @@ -114,7 +118,7 @@ def test_double_alias(self, test_name="double_alias"): output1._attrs["is_output"] = True output1._attrs["name"] = "output_1" - compile_model( + model = compile_model( [output0, output1], detect_target(), self.WORKDIR, @@ -135,11 +139,26 @@ def test_double_alias(self, test_name="double_alias"): expected_codegen = ( set_value("output_0", view_name), set_value("output_1", view_name), + device_copy(output0, view, dst_idx=1), ) self._assert_codegen_exists( test_name, expected_codegen, self.MODEL_GENERATED_FILE ) + # This is an edge case -- test the accuracy. + x_pt = get_random_torch_tensor(self.SHAPE) + gelu_pt = torch.nn.functional.gelu(x_pt) + output0_pt = torch.unsqueeze(gelu_pt, dim=0) + output1_pt = torch.flatten(gelu_pt) + output0_ait = torch.empty_like(output0_pt) + output1_ait = torch.empty_like(output1_pt) + + model.run_with_tensors( + {"input_x": x_pt}, {"output_0": output0_ait, "output_1": output1_ait} + ) + self.assertTrue(torch.allclose(output0_ait, output0_pt, atol=1e-2, rtol=1e-2)) + self.assertTrue(torch.allclose(output1_ait, output1_pt, atol=1e-2, rtol=1e-2)) + def test_output_is_view_of_output(self, test_name="output_is_view_of_output"): """ Case: An output is a view of an output.