Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!47712 add_sublist_mode_to_einsum_master 

Merge pull request !47712 from yide12/tensor_view_einsum_master
This commit is contained in:
i-robot 2023-02-13 02:33:54 +00:00 committed by Gitee
parent 1abda3f03b 3a5e32ff0d
commit bd5e820ff5
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
3 changed files with 69 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docs/api/api_python/ops/mindspore.ops.func_einsum.rst
View File

@ -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)之间。

62
mindspore/python/mindspore/ops/function/math_func.py
View File

@ -9122,12 +9122,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.
@ -9141,6 +9159,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``
@ -9151,51 +9170,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)

15
tests/st/ops/test_ops_einsum.py
View File

@ -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)