From 5abe6eed6810571cf40e2dfb0c73051d1fa5206e Mon Sep 17 00:00:00 2001 From: Akhil Goel Date: Wed, 18 Sep 2024 17:34:03 -0700 Subject: [PATCH 1/4] Add outer product implementation Signed-off-by: Akhil Goel --- tripy/tripy/frontend/ops/outer.py | 54 +++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 tripy/tripy/frontend/ops/outer.py diff --git a/tripy/tripy/frontend/ops/outer.py b/tripy/tripy/frontend/ops/outer.py new file mode 100644 index 000000000..727ed5fd0 --- /dev/null +++ b/tripy/tripy/frontend/ops/outer.py @@ -0,0 +1,54 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 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. +# + +from tripy import export, constraints +import tripy.frontend.utils as frontend_utils + + +@export.public_api(document_under="operations/functions") +@frontend_utils.convert_inputs_to_tensors(sync_arg_types=[("vec1", "vec2")]) +@constraints.dtype_info( + dtype_variables={"T1": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "int64", "bool"]}, + dtype_constraints={"vec1": "T1", "vec2": "T1", "other": "T1", constraints.RETURN_VALUE: "T1"}, +) +def outer(vec1: "tripy.Tensor", vec2: "tripy.Tensor") -> "tripy.Tensor": + r""" + Computes the outer product of 1-d vectors `vec1` and `vec2`, such that the + output dimension is (m x n) if the inputs are of size m and n respectively. + + Args: + vec1: The first 1d input vector. + vec2: The second 1d input vector + + Returns: + The outer product of the input vectors. + + .. code-block:: python + :linenos: + :caption: Example + + v1 = tp.arange(5, dtype=tp.float32) + v2 = tp.arange(4, dtype=tp.float32) + output = tp.outer(v1, v2) + + t1 = torch.arange(5, dtype=torch.float32) # doc: omit + t2 = torch.arange(4, dtype=torch.float32) # doc: omit + assert tp.allclose(output, tp.Tensor(torch.outer(t1, t2))) + """ + from tripy.frontend.trace.ops.unsqueeze import unsqueeze + + return unsqueeze(vec1, -1) * unsqueeze(vec2, 0) \ No newline at end of file From 67b53951841ab84acc75a9c7f767d6b43f0c1bc3 Mon Sep 17 00:00:00 2001 From: Akhil Goel Date: Thu, 19 Sep 2024 15:21:41 -0700 Subject: [PATCH 2/4] Add rank verification, address comments Signed-off-by: Akhil Goel --- tripy/tests/frontend/ops/test_outer.py | 23 +++++++++++++++++++++++ tripy/tripy/frontend/ops/outer.py | 16 +++++++++++----- 2 files changed, 34 insertions(+), 5 deletions(-) create mode 100644 tripy/tests/frontend/ops/test_outer.py diff --git a/tripy/tests/frontend/ops/test_outer.py b/tripy/tests/frontend/ops/test_outer.py new file mode 100644 index 000000000..17d2fd7be --- /dev/null +++ b/tripy/tests/frontend/ops/test_outer.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +from tests import helper +import tripy as tp +class TestOuter: + def test_invalid_rank_fails(self): + a = tp.ones((5, 1)) + b = tp.ones((1, 4)) + with helper.raises(tp.TripyException, "Expected input vectors to be 1-d."): + tp.outer(a, b) \ No newline at end of file diff --git a/tripy/tripy/frontend/ops/outer.py b/tripy/tripy/frontend/ops/outer.py index 727ed5fd0..fd620013e 100644 --- a/tripy/tripy/frontend/ops/outer.py +++ b/tripy/tripy/frontend/ops/outer.py @@ -20,15 +20,14 @@ @export.public_api(document_under="operations/functions") -@frontend_utils.convert_inputs_to_tensors(sync_arg_types=[("vec1", "vec2")]) @constraints.dtype_info( dtype_variables={"T1": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "int64", "bool"]}, - dtype_constraints={"vec1": "T1", "vec2": "T1", "other": "T1", constraints.RETURN_VALUE: "T1"}, + dtype_constraints={"vec1": "T1", "vec2": "T1", constraints.RETURN_VALUE: "T1"}, ) def outer(vec1: "tripy.Tensor", vec2: "tripy.Tensor") -> "tripy.Tensor": r""" - Computes the outer product of 1-d vectors `vec1` and `vec2`, such that the - output dimension is (m x n) if the inputs are of size m and n respectively. + Computes the outer product of 1-d vectors ``vec1`` and ``vec2``, such that the + output dimension is :math:`(m \times n)` if the inputs are of size :math:`m` and :math:`n` respectively. Args: vec1: The first 1d input vector. @@ -50,5 +49,12 @@ def outer(vec1: "tripy.Tensor", vec2: "tripy.Tensor") -> "tripy.Tensor": assert tp.allclose(output, tp.Tensor(torch.outer(t1, t2))) """ from tripy.frontend.trace.ops.unsqueeze import unsqueeze + from tripy.common.exception import raise_error - return unsqueeze(vec1, -1) * unsqueeze(vec2, 0) \ No newline at end of file + if vec1.rank != 1 or vec2.rank != 1: + raise_error( + "Expected input vectors to be 1-d.", + [f"Got vec1.rank={vec1.rank}, ", f"vec2.rank={vec2.rank}"], + ) + + return unsqueeze(vec1, -1) @ unsqueeze(vec2, 0) \ No newline at end of file From 187dfe71825b83ffde149932d651f8dfa6cf4c8e Mon Sep 17 00:00:00 2001 From: Akhil Goel Date: Thu, 19 Sep 2024 16:58:16 -0700 Subject: [PATCH 3/4] Docstring and example clean-up Signed-off-by: Akhil Goel --- tripy/tripy/frontend/ops/outer.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tripy/tripy/frontend/ops/outer.py b/tripy/tripy/frontend/ops/outer.py index fd620013e..d79083970 100644 --- a/tripy/tripy/frontend/ops/outer.py +++ b/tripy/tripy/frontend/ops/outer.py @@ -27,11 +27,11 @@ def outer(vec1: "tripy.Tensor", vec2: "tripy.Tensor") -> "tripy.Tensor": r""" Computes the outer product of 1-d vectors ``vec1`` and ``vec2``, such that the - output dimension is :math:`(m \times n)` if the inputs are of size :math:`m` and :math:`n` respectively. + output shape is :math:`(m, n)` if the inputs are of size :math:`(m,)` and :math:`(n,)` respectively. Args: vec1: The first 1d input vector. - vec2: The second 1d input vector + vec2: The second 1d input vector. Returns: The outer product of the input vectors. @@ -46,7 +46,9 @@ def outer(vec1: "tripy.Tensor", vec2: "tripy.Tensor") -> "tripy.Tensor": t1 = torch.arange(5, dtype=torch.float32) # doc: omit t2 = torch.arange(4, dtype=torch.float32) # doc: omit - assert tp.allclose(output, tp.Tensor(torch.outer(t1, t2))) + torch_out = torch.outer(t1, t2) # doc: omit + assert tp.allclose(output, tp.Tensor(torch_out)) + assert output.shape == torch_out.shape """ from tripy.frontend.trace.ops.unsqueeze import unsqueeze from tripy.common.exception import raise_error @@ -57,4 +59,4 @@ def outer(vec1: "tripy.Tensor", vec2: "tripy.Tensor") -> "tripy.Tensor": [f"Got vec1.rank={vec1.rank}, ", f"vec2.rank={vec2.rank}"], ) - return unsqueeze(vec1, -1) @ unsqueeze(vec2, 0) \ No newline at end of file + return unsqueeze(vec1, -1) @ unsqueeze(vec2, 0) From 01d7b6ba1677cc40ca45c40510bc27611386cb9f Mon Sep 17 00:00:00 2001 From: Akhil Goel Date: Fri, 20 Sep 2024 15:26:40 -0700 Subject: [PATCH 4/4] Add integration tests Signed-off-by: Akhil Goel --- tripy/tests/integration/test_outer.py | 38 +++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 tripy/tests/integration/test_outer.py diff --git a/tripy/tests/integration/test_outer.py b/tripy/tests/integration/test_outer.py new file mode 100644 index 000000000..91de8537b --- /dev/null +++ b/tripy/tests/integration/test_outer.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 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 pytest + +import tripy as tp +import torch + + +class TestOuter: + def test_outer(self): + v1 = tp.arange(5, dtype=tp.float32) + v2 = tp.arange(4, dtype=tp.float32) + output = tp.outer(v1, v2) + + t1 = torch.arange(5, dtype=torch.float32) + t2 = torch.arange(4, dtype=torch.float32) + torch_out = torch.outer(t1, t2) + assert tp.allclose(output, tp.Tensor(torch_out)) + assert output.shape == torch_out.shape + + def test_empty(self): + v1 = tp.Tensor([]) + v2 = tp.arange(3, dtype=tp.float32) + output = tp.outer(v1, v2) + + assert output.shape == (0, 3)