From 3d6ef1919911edf4c69537eafc4cff4196f4dcf5 Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Fri, 16 Aug 2024 11:49:19 -0700 Subject: [PATCH 01/11] added integration test for layernorm --- tripy/tests/integration/test_layernorm.py | 62 +++++++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 tripy/tests/integration/test_layernorm.py diff --git a/tripy/tests/integration/test_layernorm.py b/tripy/tests/integration/test_layernorm.py new file mode 100644 index 000000000..bb6059855 --- /dev/null +++ b/tripy/tests/integration/test_layernorm.py @@ -0,0 +1,62 @@ +# +# 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 torch +import pytest + +import tripy as tp + +DTYPES = [ + (torch.float16, tp.float16), + (torch.float32, tp.float32), +] + +@pytest.mark.parametrize("torch_dtype, tp_dtype", DTYPES) +@pytest.mark.parametrize("device", ["gpu", "cpu"]) +class TestLayerNorm: + @pytest.mark.parametrize("input_shape", [(10, 10, 5)]) + @pytest.mark.parametrize("normalized_shape", [(10, 5), (5,)]) + @pytest.mark.parametrize("eps", [1e-5, 1e-3]) + def test_layernorm_module(self, torch_dtype, tp_dtype, device, input_shape, normalized_shape, eps): + torch_device = "cuda" if device == "gpu" else device + tp_layernorm = tp.LayerNorm( + normalized_shape=normalized_shape, + eps=eps, + dtype=tp_dtype, + ) + layernorm = torch.nn.LayerNorm( + normalized_shape=normalized_shape, + eps=eps, + dtype=torch_dtype, + device=torch_device + ) + + # use Tripy's parameters + layernorm.weight = torch.nn.Parameter(torch.from_dlpack(tp_layernorm.weight.__dlpack__()).to(torch_device)) + layernorm.bias = torch.nn.Parameter(torch.from_dlpack(tp_layernorm.bias.__dlpack__()).to(torch_device)) + + input = torch.arange(torch.prod(torch.Tensor(input_shape))).reshape(input_shape).to(torch_device).to(torch_dtype) + tp_input = tp.Tensor(input, dtype=tp_dtype, device=tp.device(device)) + + output = tp_layernorm(tp_input) + expected = layernorm(input).to("cpu") + + output_torch = torch.from_dlpack(output).to("cpu") + rtol_ = 2e-7 if tp_dtype == tp.float32 else 1.5e-3 + assert torch.allclose(output_torch, expected, rtol=rtol_) + assert output_torch.shape == expected.shape \ No newline at end of file From 3a167d3c860cb77cfb9a9afa0a9afecb83f759a3 Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Tue, 20 Aug 2024 11:37:33 -0700 Subject: [PATCH 02/11] Make normalized_shape API consistent with torch --- tripy/tripy/frontend/module/layernorm.py | 30 ++++++++++++++++++++---- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/tripy/tripy/frontend/module/layernorm.py b/tripy/tripy/frontend/module/layernorm.py index 1ba892c86..3a61e6e8d 100644 --- a/tripy/tripy/frontend/module/layernorm.py +++ b/tripy/tripy/frontend/module/layernorm.py @@ -16,6 +16,7 @@ # from dataclasses import dataclass +from typing import Union, Tuple from tripy import export, utils from tripy.common import datatype @@ -33,11 +34,17 @@ class LayerNorm(Module): :math:`\text{LayerNorm}(x) = \Large \frac{x - \bar{x}}{ \sqrt{\sigma^2 + \epsilon}} \normalsize * \gamma + \beta` where :math:`\bar{x}` is the mean and :math:`\sigma^2` is the variance. + + The mean the standard deviation are calculated over the last :math:`D` + dimensions, where :math:`D` is the dimension of `normalized_shape`. """ dtype: datatype.dtype r"""The data type used to perform the operation.""" + normalized_shape: Tuple[int] + r"""Defines the shape of the input tensor that is to be normalized over.""" + weight: Parameter r"""The :math:`\gamma` parameter of shape :math:`[\text{normalized_shape}]`.""" @@ -47,7 +54,7 @@ class LayerNorm(Module): eps: float """A value added to the denominator to prevent division by zero.""" - def __init__(self, normalized_shape: int, dtype: datatype.dtype = datatype.float32, eps: float = 1e-5) -> None: + def __init__(self, normalized_shape: Union[int, Tuple[int]], dtype: datatype.dtype = datatype.float32, eps: float = 1e-5) -> None: """ Args: normalized_shape: The size of the feature dimension of the input over which normalization is performed. @@ -77,9 +84,14 @@ def __init__(self, normalized_shape: int, dtype: datatype.dtype = datatype.float self.dtype = dtype # Replace with random weights when #74 is completed. - self.weight = DefaultParameter((normalized_shape,), dtype=dtype) + if isinstance(normalized_shape, int): + normalized_shape = (normalized_shape,) + + self.normalized_shape = normalized_shape - self.bias = DefaultParameter((normalized_shape,), dtype=dtype) + self.weight = DefaultParameter(normalized_shape, dtype=dtype) + + self.bias = DefaultParameter(normalized_shape, dtype=dtype) self.eps = eps @@ -93,8 +105,16 @@ def __call__(self, x: "tripy.Tensor") -> "tripy.Tensor": """ from tripy.frontend.trace.ops.reduce import mean, var from tripy.frontend.trace.ops.unary_elementwise import rsqrt + from tripy.common.exception import raise_error + import tripy as tp + + # The mean and the variance are computed over the last D dimensions + D = len(self.normalized_shape) + + if x[-D:].shape != tp.Shape(self.normalized_shape): + raise_error(f"The input's last {D} dimensions must have a shape of {self.normalized_shape}") - mean_val = mean(x, dim=-1, keepdim=True) - var_val = var(x, dim=-1, keepdim=True, correction=0) + self.eps + mean_val = mean(x, dim=-D, keepdim=True) + var_val = var(x, dim=-D, keepdim=True, correction=0) + self.eps x = (x - mean_val) * rsqrt(var_val) return self.weight * x + self.bias From e1598fa99ad9e6d0680c1938880ce4c89902ed9b Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Thu, 22 Aug 2024 12:05:09 -0700 Subject: [PATCH 03/11] Add correction as a param; fix error msg for incorrect dims --- tripy/tripy/frontend/module/layernorm.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tripy/tripy/frontend/module/layernorm.py b/tripy/tripy/frontend/module/layernorm.py index 3a61e6e8d..de4603eef 100644 --- a/tripy/tripy/frontend/module/layernorm.py +++ b/tripy/tripy/frontend/module/layernorm.py @@ -54,7 +54,10 @@ class LayerNorm(Module): eps: float """A value added to the denominator to prevent division by zero.""" - def __init__(self, normalized_shape: Union[int, Tuple[int]], dtype: datatype.dtype = datatype.float32, eps: float = 1e-5) -> None: + correction: int + """Difference between the sample size and the degrees of freedom.""" + + def __init__(self, normalized_shape: Union[int, Tuple[int]], dtype: datatype.dtype = datatype.float32, eps: float = 1e-5, correction: int = 0) -> None: """ Args: normalized_shape: The size of the feature dimension of the input over which normalization is performed. @@ -95,6 +98,8 @@ def __init__(self, normalized_shape: Union[int, Tuple[int]], dtype: datatype.dty self.eps = eps + self.correction = correction + def __call__(self, x: "tripy.Tensor") -> "tripy.Tensor": r""" Args: @@ -111,10 +116,11 @@ def __call__(self, x: "tripy.Tensor") -> "tripy.Tensor": # The mean and the variance are computed over the last D dimensions D = len(self.normalized_shape) - if x[-D:].shape != tp.Shape(self.normalized_shape): - raise_error(f"The input's last {D} dimensions must have a shape of {self.normalized_shape}") + if x.shape[-D:] != tp.Shape(self.normalized_shape): + raise_error(f"The input's last {D} dimensions must have a shape of {self.normalized_shape} and received {x.shape[-D:].as_tensor().data()}") - mean_val = mean(x, dim=-D, keepdim=True) - var_val = var(x, dim=-D, keepdim=True, correction=0) + self.eps + reduce_dims = tuple(-i for i in range(D, 0, -1)) + mean_val = mean(x, dim=reduce_dims, keepdim=True) + var_val = var(x, dim=reduce_dims, keepdim=True, correction=self.correction) + self.eps x = (x - mean_val) * rsqrt(var_val) return self.weight * x + self.bias From 7d696c06804850928cc520f9abc6ba7c97b6deb0 Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Thu, 22 Aug 2024 12:05:50 -0700 Subject: [PATCH 04/11] fix integration test; add test for incorrect dims --- tripy/tests/integration/test_layernorm.py | 48 +++++++++++++---------- 1 file changed, 27 insertions(+), 21 deletions(-) diff --git a/tripy/tests/integration/test_layernorm.py b/tripy/tests/integration/test_layernorm.py index bb6059855..7f297fb49 100644 --- a/tripy/tests/integration/test_layernorm.py +++ b/tripy/tests/integration/test_layernorm.py @@ -16,47 +16,53 @@ # import numpy as np +import re import torch import pytest import tripy as tp +from tripy.common.exception import TripyException DTYPES = [ (torch.float16, tp.float16), - (torch.float32, tp.float32), + (torch.float32, tp.float32) ] -@pytest.mark.parametrize("torch_dtype, tp_dtype", DTYPES) -@pytest.mark.parametrize("device", ["gpu", "cpu"]) class TestLayerNorm: - @pytest.mark.parametrize("input_shape", [(10, 10, 5)]) - @pytest.mark.parametrize("normalized_shape", [(10, 5), (5,)]) + @pytest.mark.parametrize("torch_dtype, tp_dtype", DTYPES) + @pytest.mark.parametrize("input_shape", [(2, 2, 2)]) + @pytest.mark.parametrize("normalized_shape", [(2, 2), (2,)]) @pytest.mark.parametrize("eps", [1e-5, 1e-3]) - def test_layernorm_module(self, torch_dtype, tp_dtype, device, input_shape, normalized_shape, eps): - torch_device = "cuda" if device == "gpu" else device - tp_layernorm = tp.LayerNorm( + def test_layernorm_accuracy(self, torch_dtype, tp_dtype, input_shape, normalized_shape, eps): + layernorm = torch.nn.LayerNorm( normalized_shape=normalized_shape, eps=eps, - dtype=tp_dtype, + dtype=torch_dtype, ) - layernorm = torch.nn.LayerNorm( + tp_layernorm = tp.LayerNorm( normalized_shape=normalized_shape, eps=eps, - dtype=torch_dtype, - device=torch_device + dtype=tp_dtype, ) # use Tripy's parameters - layernorm.weight = torch.nn.Parameter(torch.from_dlpack(tp_layernorm.weight.__dlpack__()).to(torch_device)) - layernorm.bias = torch.nn.Parameter(torch.from_dlpack(tp_layernorm.bias.__dlpack__()).to(torch_device)) + tp_layernorm.weight = tp.Parameter(layernorm.weight) + tp_layernorm.bias = tp.Parameter(layernorm.bias) - input = torch.arange(torch.prod(torch.Tensor(input_shape))).reshape(input_shape).to(torch_device).to(torch_dtype) - tp_input = tp.Tensor(input, dtype=tp_dtype, device=tp.device(device)) + input = torch.arange(torch.prod(torch.Tensor(input_shape))).reshape(input_shape).to(torch_dtype) + tp_input = tp.Tensor(input, dtype=tp_dtype) output = tp_layernorm(tp_input) - expected = layernorm(input).to("cpu") + expected = tp.Tensor(layernorm(input), device=tp.device("cpu")) + + rtol_ = 2e-7 if tp_dtype == tp.float32 else 1e-3 + assert output.shape == expected.shape + assert tp.allclose(output, expected, rtol=rtol_) - output_torch = torch.from_dlpack(output).to("cpu") - rtol_ = 2e-7 if tp_dtype == tp.float32 else 1.5e-3 - assert torch.allclose(output_torch, expected, rtol=rtol_) - assert output_torch.shape == expected.shape \ No newline at end of file + def test_layernorm_improper_dimensions(self): + tp_layernorm = tp.LayerNorm( + normalized_shape=[2, 2], + ) + x = tp.ones((5,5,5)) + with pytest.raises(TripyException, match=re.escape("The input's last 2 dimensions must have a shape of [2, 2] and received [5, 5]")): + tp_layernorm(x) \ No newline at end of file From c28b56aac9601768b14763b0f07a83f51d4425b8 Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Thu, 22 Aug 2024 12:14:07 -0700 Subject: [PATCH 05/11] Use Shape directly --- tripy/tripy/frontend/module/layernorm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tripy/tripy/frontend/module/layernorm.py b/tripy/tripy/frontend/module/layernorm.py index de4603eef..12ac84df1 100644 --- a/tripy/tripy/frontend/module/layernorm.py +++ b/tripy/tripy/frontend/module/layernorm.py @@ -109,14 +109,14 @@ def __call__(self, x: "tripy.Tensor") -> "tripy.Tensor": A tensor of the same shape as the input. """ from tripy.frontend.trace.ops.reduce import mean, var + from tripy.frontend.shape import Shape from tripy.frontend.trace.ops.unary_elementwise import rsqrt from tripy.common.exception import raise_error - import tripy as tp # The mean and the variance are computed over the last D dimensions D = len(self.normalized_shape) - if x.shape[-D:] != tp.Shape(self.normalized_shape): + if x.shape[-D:] != Shape(self.normalized_shape): raise_error(f"The input's last {D} dimensions must have a shape of {self.normalized_shape} and received {x.shape[-D:].as_tensor().data()}") reduce_dims = tuple(-i for i in range(D, 0, -1)) From 3c8a935ca18283c44cab0ba1c26d1bbe5c26d836 Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Thu, 22 Aug 2024 16:27:34 -0700 Subject: [PATCH 06/11] Remove correction as parameter, default to 0; add normalized_shape to constant fields --- tripy/tripy/frontend/module/layernorm.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/tripy/tripy/frontend/module/layernorm.py b/tripy/tripy/frontend/module/layernorm.py index 12ac84df1..fb1a38bd3 100644 --- a/tripy/tripy/frontend/module/layernorm.py +++ b/tripy/tripy/frontend/module/layernorm.py @@ -26,7 +26,7 @@ @export.public_api(document_under="operations/modules") @dataclass -@utils.constant_fields(["dtype"]) +@utils.constant_fields(["dtype", "normalized_shape"]) class LayerNorm(Module): r""" Applies layer normalization over the input tensor: @@ -54,10 +54,7 @@ class LayerNorm(Module): eps: float """A value added to the denominator to prevent division by zero.""" - correction: int - """Difference between the sample size and the degrees of freedom.""" - - def __init__(self, normalized_shape: Union[int, Tuple[int]], dtype: datatype.dtype = datatype.float32, eps: float = 1e-5, correction: int = 0) -> None: + def __init__(self, normalized_shape: Union[int, Tuple[int]], dtype: datatype.dtype = datatype.float32, eps: float = 1e-5) -> None: """ Args: normalized_shape: The size of the feature dimension of the input over which normalization is performed. @@ -98,8 +95,6 @@ def __init__(self, normalized_shape: Union[int, Tuple[int]], dtype: datatype.dty self.eps = eps - self.correction = correction - def __call__(self, x: "tripy.Tensor") -> "tripy.Tensor": r""" Args: @@ -121,6 +116,6 @@ def __call__(self, x: "tripy.Tensor") -> "tripy.Tensor": reduce_dims = tuple(-i for i in range(D, 0, -1)) mean_val = mean(x, dim=reduce_dims, keepdim=True) - var_val = var(x, dim=reduce_dims, keepdim=True, correction=self.correction) + self.eps + var_val = var(x, dim=reduce_dims, keepdim=True, correction=0) + self.eps x = (x - mean_val) * rsqrt(var_val) return self.weight * x + self.bias From ea19442596edfb8c808c17b5a2ec569b2be6ef2a Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Thu, 22 Aug 2024 16:30:30 -0700 Subject: [PATCH 07/11] Add eps to groupnorm --- tripy/tripy/frontend/module/groupnorm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tripy/tripy/frontend/module/groupnorm.py b/tripy/tripy/frontend/module/groupnorm.py index 80150963a..e766be8fa 100644 --- a/tripy/tripy/frontend/module/groupnorm.py +++ b/tripy/tripy/frontend/module/groupnorm.py @@ -55,7 +55,7 @@ class GroupNorm(Module): eps: float """A value added to the denominator to prevent division by zero. Defaults to 1e-5.""" - def __init__(self, num_groups: int, num_channels: int, dtype: datatype.dtype = datatype.float32) -> None: + def __init__(self, num_groups: int, num_channels: int, dtype: datatype.dtype = datatype.float32, eps: float = 1e-5) -> None: """ Args: num_groups: The number of groups to split the channels into. @@ -97,7 +97,7 @@ def __init__(self, num_groups: int, num_channels: int, dtype: datatype.dtype = d # Replace with random weights when #74 is completed. self.weight = DefaultParameter((num_channels,), dtype=dtype) self.bias = DefaultParameter((num_channels,), dtype=dtype) - self.eps = 1e-5 + self.eps = eps def __call__(self, x: "tripy.Tensor") -> "tripy.Tensor": r""" From 99591a3fb8eaf9f3fc499fbbcddd9490b09b2010 Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Thu, 22 Aug 2024 16:34:34 -0700 Subject: [PATCH 08/11] Add groupnorm integration test --- tripy/tests/integration/test_groupnorm.py | 61 +++++++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 tripy/tests/integration/test_groupnorm.py diff --git a/tripy/tests/integration/test_groupnorm.py b/tripy/tests/integration/test_groupnorm.py new file mode 100644 index 000000000..0b8bd7c80 --- /dev/null +++ b/tripy/tests/integration/test_groupnorm.py @@ -0,0 +1,61 @@ +# +# 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 torch +import pytest + +import tripy as tp +from tripy.common.exception import TripyException + +DTYPES = [ + (torch.float16, tp.float16), + (torch.float32, tp.float32) +] + +class TestGroupNorm: + @pytest.mark.parametrize("torch_dtype, tp_dtype", DTYPES) + @pytest.mark.parametrize("input_shape", [(2, 2, 2)]) + @pytest.mark.parametrize("num_groups", [2, 4]) + @pytest.mark.parametrize("num_channels", [2]) + @pytest.mark.parametrize("eps", [1e-5, 1e-3]) + def test_groupnorm_accuracy(self, torch_dtype, tp_dtype, input_shape, num_groups, num_channels, eps): + groupnorm = torch.nn.GroupNorm( + num_groups=num_groups, + num_channels=num_channels, + eps=eps, + dtype=torch_dtype, + ) + tp_groupnorm = tp.GroupNorm( + num_groups=num_groups, + num_channels=num_channels, + eps=eps, + dtype=tp_dtype, + ) + + tp_groupnorm.weight = tp.Parameter(groupnorm.weight) + tp_groupnorm.bias = tp.Parameter(groupnorm.bias) + + input = torch.arange(torch.prod(torch.Tensor(input_shape))).reshape(input_shape).to(torch_dtype) + tp_input = tp.Tensor(input, dtype=tp_dtype) + + output = tp_groupnorm(tp_input) + expected = tp.Tensor(groupnorm(input), device=tp.device("cpu")) + + rtol_ = 2e-6 if tp_dtype == tp.float32 else 1e-3 + assert output.shape == expected.shape + assert tp.allclose(output, expected, rtol=rtol_) \ No newline at end of file From 59a8b3c42437587040428082636031d6b1a32037 Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Thu, 22 Aug 2024 16:36:37 -0700 Subject: [PATCH 09/11] Fix groupnorm integration test --- tripy/tests/integration/test_groupnorm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tripy/tests/integration/test_groupnorm.py b/tripy/tests/integration/test_groupnorm.py index 0b8bd7c80..11ac2967e 100644 --- a/tripy/tests/integration/test_groupnorm.py +++ b/tripy/tests/integration/test_groupnorm.py @@ -29,9 +29,9 @@ class TestGroupNorm: @pytest.mark.parametrize("torch_dtype, tp_dtype", DTYPES) - @pytest.mark.parametrize("input_shape", [(2, 2, 2)]) - @pytest.mark.parametrize("num_groups", [2, 4]) - @pytest.mark.parametrize("num_channels", [2]) + @pytest.mark.parametrize("input_shape", [(1, 10, 2)]) + @pytest.mark.parametrize("num_groups", [2, 5]) + @pytest.mark.parametrize("num_channels", [10]) @pytest.mark.parametrize("eps", [1e-5, 1e-3]) def test_groupnorm_accuracy(self, torch_dtype, tp_dtype, input_shape, num_groups, num_channels, eps): groupnorm = torch.nn.GroupNorm( From 132c709efbadac9d992c34c0ee4d0b18e94d3e45 Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Fri, 23 Aug 2024 08:34:26 -0700 Subject: [PATCH 10/11] fix typo in docstring; remove explicit Shape construction --- tripy/tripy/frontend/module/layernorm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tripy/tripy/frontend/module/layernorm.py b/tripy/tripy/frontend/module/layernorm.py index fb1a38bd3..a54fce6e1 100644 --- a/tripy/tripy/frontend/module/layernorm.py +++ b/tripy/tripy/frontend/module/layernorm.py @@ -35,7 +35,7 @@ class LayerNorm(Module): where :math:`\bar{x}` is the mean and :math:`\sigma^2` is the variance. - The mean the standard deviation are calculated over the last :math:`D` + The mean and standard deviation are calculated over the last :math:`D` dimensions, where :math:`D` is the dimension of `normalized_shape`. """ @@ -111,7 +111,7 @@ def __call__(self, x: "tripy.Tensor") -> "tripy.Tensor": # The mean and the variance are computed over the last D dimensions D = len(self.normalized_shape) - if x.shape[-D:] != Shape(self.normalized_shape): + if x.shape[-D:] != self.normalized_shape: raise_error(f"The input's last {D} dimensions must have a shape of {self.normalized_shape} and received {x.shape[-D:].as_tensor().data()}") reduce_dims = tuple(-i for i in range(D, 0, -1)) From 27a6ee882edeb6285c31bab054770886dc9b8da6 Mon Sep 17 00:00:00 2001 From: Mark Kraay Date: Fri, 23 Aug 2024 08:43:28 -0700 Subject: [PATCH 11/11] fix LayerNorm error msg --- tripy/tripy/frontend/module/layernorm.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tripy/tripy/frontend/module/layernorm.py b/tripy/tripy/frontend/module/layernorm.py index a54fce6e1..d62c4bb33 100644 --- a/tripy/tripy/frontend/module/layernorm.py +++ b/tripy/tripy/frontend/module/layernorm.py @@ -112,7 +112,9 @@ def __call__(self, x: "tripy.Tensor") -> "tripy.Tensor": D = len(self.normalized_shape) if x.shape[-D:] != self.normalized_shape: - raise_error(f"The input's last {D} dimensions must have a shape of {self.normalized_shape} and received {x.shape[-D:].as_tensor().data()}") + raise_error("Unexpected input shape", + [f"The input's last {D} dimensions must have a shape of {self.normalized_shape} and received {x.shape[-D:].data()}"] + ) reduce_dims = tuple(-i for i in range(D, 0, -1)) mean_val = mean(x, dim=reduce_dims, keepdim=True)