From 3a5e32ff0db6686463dd215decba53431f85494f Mon Sep 17 00:00:00 2001 From: yide12 Date: Tue, 10 Jan 2023 15:49:08 +0800 Subject: [PATCH] add_sublist_mode_to_einsum_master --- .../ops/mindspore.ops.func_einsum.rst | 4 ++ .../mindspore/ops/function/math_func.py | 62 +++++++++++++++---- tests/st/ops/test_ops_einsum.py | 15 +++++ 3 files changed, 69 insertions(+), 12 deletions(-) diff --git a/docs/api/api_python/ops/mindspore.ops.func_einsum.rst b/docs/api/api_python/ops/mindspore.ops.func_einsum.rst index c11b43324a2..ed6dfb9aaf4 100644 --- a/docs/api/api_python/ops/mindspore.ops.func_einsum.rst +++ b/docs/api/api_python/ops/mindspore.ops.func_einsum.rst @@ -5,6 +5,9 @@ mindspore.ops.einsum 基于爱因斯坦求和约定（Einsum）符号，指定维度对输入Tensor元素的乘积求和。你可以使用这个运算符来执行对角线、减法、转置、矩阵乘法、乘法、内积运算等等。 + .. note:: + 现在支持子列表模式。例如，ops.einsum(op1, sublist1, op2, sublist2, ..., sublist_out)。在子列表模式中，`equation` 由子列表推导得到，Python的省略号和介于[0, 52)的整数list组成子列表。每个操作数后面都有一个子列表，并且最后有一个表示输出的子列表。 + 参数： - **equation** (str) - 基于爱因斯坦求和约定的符号，表示想要执行的操作。符号只能包含字母、逗号、省略号和箭头。字母表示输入Tensor维数，逗号表示单独的Tensor，省略号表示忽略的Tensor维数，箭头的左边表示输入Tensor，右边表示期望输出的维度。 - **operands** (Tensor) - 用于计算的输入Tensor。Tensor的数据类型必须相同。 @@ -14,3 +17,4 @@ mindspore.ops.einsum 异常： - **TypeError** - `equation` 无效或者不匹配输入Tensor。 + - **ValueError** - 子列表模式下子列表的数字不介于[0, 52)之间。 diff --git a/mindspore/python/mindspore/ops/function/math_func.py b/mindspore/python/mindspore/ops/function/math_func.py index 5dd6b81f618..1051fb76327 100644 --- a/mindspore/python/mindspore/ops/function/math_func.py +++ b/mindspore/python/mindspore/ops/function/math_func.py @@ -8754,12 +8754,30 @@ def cross(input, other, dim=None): return cross_op(input, other) +def _einsum_convert_num_to_char(num): + """For einsum, convert number into char.""" + if [num] == [Ellipsis]: + return '...' + # pylint: disable=chained-comparison + if num >= 0 and num < 26: + return chr(num + ord('A')) + # pylint: disable=chained-comparison + if num >= 26 and num < 52: + return chr(num - 26 + ord('a')) + raise ValueError(f"For Einsum, the number in sublist should be in range [0， 52), but got {num}") + + def einsum(equation, *operands): r""" Sums the product of the elements of the input Tensor along dimensions specified notation based on the Einstein summation convention(Einsum). You can use this operator to perform diagonal, reducesum, transpose, matmul, mul, inner product operations, etc. + Note:: + The sublist format is alse supported. For example, ops.einsum(op1, sublist1, op2, sublist2, ..., sublist_out). + In this format, equation can be derived by the sublists which are made up of Python's Ellipsis and list of + integers in [0, 52). Each operand is followed by a sublist and an output sublist is at the end. + Args: equation (str): Notation based on the Einstein summation convention, represent the operation you want to do. the value can contain only letters, commas, ellipsis and arrow. @@ -8773,6 +8791,7 @@ def einsum(equation, *operands): Raises: TypeError: If `equation` is invalid, or the `equation` does not match the input tensor. + ValueError: If the number in sublist is not in [0, 52) in sublist format. Supported Platforms: ``GPU`` @@ -8783,51 +8802,70 @@ def einsum(equation, *operands): >>> output = ops.einsum(equation, x) >>> print(output) [7.] - >>> >>> x = Tensor(np.array([1.0, 2.0, 4.0]), mindspore.float32) >>> y = Tensor(np.array([2.0, 4.0, 3.0]), mindspore.float32) >>> equation = "i,i->i" >>> output = ops.einsum(equation, x, y) >>> print(output) [ 2. 8. 12.] - >>> >>> x = Tensor(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), mindspore.float32) >>> y = Tensor(np.array([[2.0, 3.0], [1.0, 2.0], [4.0, 5.0]]), mindspore.float32) >>> equation = "ij,jk->ik" >>> output = ops.einsum(equation, x, y) >>> print(output) [[16. 22.] - [37. 52.]] - >>> + [37. 52.]] >>> x = Tensor(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), mindspore.float32) >>> equation = "ij->ji" >>> output = ops.einsum(equation, x) >>> print(output) [[1. 4.] - [2. 5.] - [3. 6.]] - >>> + [2. 5.] + [3. 6.]] >>> x = Tensor(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), mindspore.float32) >>> equation = "ij->j" >>> output = ops.einsum(equation, x) >>> print(output) [5. 7. 9.] - >>> >>> x = Tensor(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), mindspore.float32) >>> equation = "...->" - >>> output = einsum(equation, x) + >>> output = ops.einsum(equation, x) >>> print(output) [21.] - >>> >>> x = Tensor(np.array([1.0, 2.0, 3.0]), mindspore.float32) >>> y = Tensor(np.array([2.0, 4.0, 1.0]), mindspore.float32) >>> equation = "j,i->ji" >>> output = ops.einsum(equation, x, y) >>> print(output) [[ 2. 4. 1.] - [ 4. 8. 2.] - [ 6. 12. 3.]] + [ 4. 8. 2.] + [ 6. 12. 3.]] + >>> x = mindspore.Tensor([1, 2, 3, 4], mindspore.float32) + >>> y = mindspore.Tensor([1, 2], mindspore.float32) + >>> output = ops.einsum(x, [..., 1], y, [..., 2], [..., 1, 2]) + [[1. 2.] + [2. 4.] + [3. 6.] + [4. 8.]] """ + if isinstance(equation, Tensor): + equ_tmp = '' + for i, lst in enumerate(operands): + if i % 2 == 0: + for _, num in enumerate(lst): + equ_tmp += _einsum_convert_num_to_char(num) + if i in (len(operands) - 1, len(operands) - 2): + continue + equ_tmp += ',' + if len(operands) % 2 == 0: + equ_tmp += '->' + for _, num in enumerate(operands[-1]): + equ_tmp += _einsum_convert_num_to_char(num) + operands_tmp = list([equation]) + list(operands[1:-1:2]) + else: + operands_tmp = list([equation]) + list(operands[1::2]) + equation = equ_tmp + operands = tuple(operands_tmp) return _get_cache_prim(P.Einsum)(equation)(operands) diff --git a/tests/st/ops/test_ops_einsum.py b/tests/st/ops/test_ops_einsum.py index 54f4e8fd751..1923db0e0f4 100644 --- a/tests/st/ops/test_ops_einsum.py +++ b/tests/st/ops/test_ops_einsum.py @@ -28,6 +28,13 @@ class Net(nn.Cell): return ops.einsum(self.equation, *operands) +class NetSublist(nn.Cell): + """Test ops.einsum in sublist format.""" + + def construct(self, x, y): + return ops.einsum(x, [..., 0, 1], y, [..., 1, 2], [..., 0, 2]) + + @pytest.mark.level0 @pytest.mark.platform_x86_gpu_training @pytest.mark.env_onecard @@ -96,3 +103,11 @@ def test_ops_einsum(mode): [4., 8., 2.], [6., 12., 3.]] assert np.allclose(output.asnumpy(), expect_output) + + x = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], ms.float32) + y = Tensor([[2.0, 3.0], [1.0, 2.0], [4.0, 5.0]], ms.float32) + net = NetSublist() + output = net(x, y) + expect_output = [[16., 22.], + [37., 52.]] + assert np.allclose(output.asnumpy(), expect_output)