Merge pull request #2 from derwind/develop

Rough implementation of `einsum`

Author: derwind

.github/workflows/tests.yml

@@ -1,9 +1,9 @@
 name: Tests
 on:
   push:
-    branches: [master, develop]
+    branches: [main, develop]
   pull_request:
-    branches: [master, develop]
+    branches: [main, develop]
 concurrency:
   group: ${{ github.repository }}-${{ github.ref }}-${{ github.head_ref }}-${{ github.workflow }}
   cancel-in-progress: true

mynumpy/core/ndarray.py

@@ -1,5 +1,5 @@
 import copy
-from typing import List, Tuple, Union, Optional, Any
+from typing import List, Tuple, Dict, Union, Optional, Any
 from ..dtypes import Numbers
 
 
@@ -57,9 +57,13 @@ def __mul__(self, other: Union[Numbers, 'ndarray']) -> 'ndarray':
 
     def __matmul__(self, other: 'ndarray') -> 'ndarray':
         if len(self.shape) < 1:
-            raise ValueError(f'matmul: Input operand 0 does not have enough dimensions (has 0, gufunc core with signature (n?,k),(k,m?)->(n?,m?) requires 1)')
+            raise ValueError(
+                f'matmul: Input operand 0 does not have enough dimensions (has 0, gufunc core with signature (n?,k),(k,m?)->(n?,m?) requires 1)'
+            )
         if len(other.shape) < 1:
-            raise ValueError(f'matmul: Input operand 1 does not have enough dimensions (has 0, gufunc core with signature (n?,k),(k,m?)->(n?,m?) requires 1)')
+            raise ValueError(
+                f'matmul: Input operand 1 does not have enough dimensions (has 0, gufunc core with signature (n?,k),(k,m?)->(n?,m?) requires 1)'
+            )
 
         if len(self.shape) != 1 and len(self.shape) != 2:
             raise ValueError(f'matmul: Input operand 0 is neither a vector nor a matrix and not supported')
@@ -81,7 +85,9 @@ def __matmul__(self, other: 'ndarray') -> 'ndarray':
             squeeze_count += 1
 
         if a.shape[1] != b.shape[0]:
-            raise ValueError(f'matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size {b.shape[0]} is different from {a.shape[1]})')
+            raise ValueError(
+                f'matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size {b.shape[0]} is different from {a.shape[1]})'
+            )
 
         n_row = a.shape[0]
         n_col = b.shape[1]
@@ -299,10 +305,118 @@ def broadcast(a, shape: Union[List[int], Tuple[int]]) -> 'ndarray':
 def einsum(subscripts: str, *operands: List[ndarray]) -> ndarray:
     subscripts = subscripts.replace(' ', '')
 
+    from_indices, to_index = subscripts.split('->')
+    if len(from_indices.split(',')) != len(operands):
+        raise ValueError('more operands provided to einstein sum function than specified in the subscripts string')
+
+    index_list = [[idx for idx in index] for index in from_indices.split(',')]
+    to_index = [idx for idx in to_index]
+
+    for i, (op, index) in enumerate(zip(operands, index_list)):
+        if len(op.shape) > len(index):
+            raise ValueError('operand has more dimensions than subscripts given in einstein sum')
+
+        if len(op.shape) < len(index):
+            raise ValueError(f'einstein sum subscripts string contains too many subscripts for operand {i}')
+
     if len(operands) != 2:
         raise ValueError(f'operands whose length != 2 are currently not supported')
 
     a, b = operands
-    from_, to_ = subscripts.split('->')
+    index_a, index_b = index_list
+
+    # index char -> loc, e.g. {'i': 0, 'j': 1, 'k': 2, 'l': 3} for 'ijkl'
+    i2l_a = {index: index_a.index(index) for index in index_a}
+    i2l_b = {index: index_b.index(index) for index in index_b}
+
+    # determin output tensor's shape
+
+    out_shape = []
+    # index char -> dim, e.g. {'i': 3, 'j': 4}
+    i2d = {}
+    for idx in to_index:
+        if idx in i2l_a:
+            dim = a.shape[i2l_a[idx]]
+            out_shape.append(dim)
+            i2d[idx] = dim
+            continue
+        if idx in i2l_b:
+            dim = b.shape[i2l_b[idx]]
+            out_shape.append(dim)
+            i2d[idx] = dim
+            continue
+        raise ValueError(f"einstein sum subscripts string included output subscript '{idx}' which never appeared in an input")
+
+    # Preprocess finished. Main process begins
+
+    placeholder = zeros(out_shape).data
+
+    def fill_placeholder(target: ndarray, index: List[str], index_kv: Optional[Dict[str, int]] = None):
+        if index_kv is None:
+            index_kv = {}
+
+        idx, index = index[0], index[1:]  # index chars
+
+        for i in range(i2d[idx]):
+            index_kv_ = index_kv.copy()
+            index_kv_[idx] = i
+            if isinstance(target[i], list):
+                fill_placeholder(target[i], index, index_kv_)
+                continue
+
+            target[i] = calc_value(a, b, index_a, index_b, index_kv_)
+
+    # e.g. 'ijkl,jmln->ikm': sum_j sum_l sum_n A_{ijkl} B_{jmln}
+    def calc_value(a_1: ndarray, a_2: ndarray, index_1: Tuple[str, ...], index_2: Tuple[str, ...], index_kv: Dict[str, int]):
+        combinations_kv = []
+        calc_combinations(list(a_1.shape), list(a_2.shape), index_1, index_2, index_kv, combinations_kv)
+
+        v = 0
+        for idx_kv in combinations_kv:
+            v_1 = get_value(a_1.data, index_1, idx_kv)
+            v_2 = get_value(a_2.data, index_2, idx_kv)
+            v += v_1 * v_2
+
+        return v
+
+    def calc_combinations(
+        shape_1: List[int], shape_2: List[int], index_1: List[str], index_2: List[str], index_kv: Dict[str, int], out_combs: List[Dict[str, int]]
+    ):
+        if index_1:
+            idx1, index_1 = index_1[0], index_1[1:]
+            dim1, shape_1 = shape_1[0], shape_1[1:]
+            if idx1 in index_kv:
+                calc_combinations(shape_1, shape_2, index_1, index_2, index_kv, out_combs)
+                return
+            else:
+                for i in range(dim1):
+                    index_kv_ = index_kv.copy()
+                    index_kv_[idx1] = i
+                    calc_combinations(shape_1, shape_2, index_1, index_2, index_kv_, out_combs)
+                return
+
+        if index_2:
+            idx2, index_2 = index_2[0], index_2[1:]
+            dim2, shape_2 = shape_2[0], shape_2[1:]
+            if idx2 in index_kv:
+                calc_combinations(shape_1, shape_2, index_1, index_2, index_kv, out_combs)
+                return
+            else:
+                for i in range(dim2):
+                    index_kv_ = index_kv.copy()
+                    index_kv_[idx2] = i
+                    calc_combinations(shape_1, shape_2, index_1, index_2, index_kv_, out_combs)
+                return
+
+        out_combs.append(index_kv)
+
+    def get_value(target: List[Numbers], index: List[Numbers], index_kv: Dict[str, int]):
+        if isinstance(target, list):
+            idx, index = index[0], index[1:]
+            target = target[index_kv[idx]]
+            return get_value(target, index, index_kv)
+        return target
+
+    fill_placeholder(placeholder, to_index)
 
-    raise NotImplementedError('not implemented yet')
+    return ndarray(placeholder)

test/ndarray_test.py

@@ -1245,3 +1245,105 @@ def test_matmul(self):
 
         with self.assertRaises(ValueError):
             b @ a
+
+    def test_einsum(self):
+        a = mynp.array([1, 2])
+
+        b = mynp.array([
+            [1, 2],
+            [3, 4]
+        ])
+
+        self.assertEqual(mynp.einsum('i,ij->j', a, b).data, [7, 10])
+        self.assertEqual(mynp.einsum('i,ji->j', a, b).data, [5, 11])
+
+        self.assertEqual(mynp.einsum('ij,i->j', b, a).data, [7, 10])
+        self.assertEqual(mynp.einsum('ij,j->i', b, a).data, [5, 11])
+
+        a = mynp.array([
+            [1, 2],
+            [3, 4]
+        ])
+
+        b = mynp.array([
+            [-2, 1],
+            [-5, 3]
+        ])
+
+        self.assertEqual(mynp.einsum('ij,jk->ik', a, b).data, [
+            [-12, 7],
+            [-26, 15]
+        ])
+
+        self.assertEqual(mynp.einsum('jk,ki->ji', a, b).data, [
+            [-12, 7],
+            [-26, 15]
+        ])
+
+        a = mynp.array([
+            [1, 2],
+            [3, 4],
+            [5, 6]
+        ])
+
+        b = mynp.array([
+            [7, 8, 9, 10],
+            [11, 12, 13, 14]
+        ])
+
+        self.assertEqual(mynp.einsum('ij,jk->ik', a, b).data, [
+            [ 29, 32, 35, 38],
+            [ 65, 72, 79, 86],
+            [101, 112, 123, 134]
+        ])
+
+        a = mynp.array([
+            [
+                [1, 2],
+                [3, 4]
+            ],
+            [
+                [5, 6],
+                [7, 8]
+            ],
+        ])
+
+        b = mynp.array([
+            [
+                [-1, -5],
+                [-3, 2],
+                [1, 4],
+            ],
+            [
+                [3, 6],
+                [-3, 2],
+                [-4, 1],
+            ]
+        ])
+
+        self.assertEqual(mynp.einsum('ijk,ilj->jl', a, b).data, [
+            [30, -42, -41],
+            [55, 44, 43]
+        ])
+
+        with self.assertRaises(ValueError):
+            data = mynp.array([
+                [1, 2],
+                [3, 4]
+            ])
+            data2 = mynp.array([
+                [-2, 1],
+                [-5, 3]
+            ])
+            mynp.einsum('ijl,jk->ik', a, b)
+
+        with self.assertRaises(ValueError):
+            data = mynp.array([
+                [1, 2],
+                [3, 4]
+            ])
+            data2 = mynp.array([
+                [-2, 1],
+                [-5, 3]
+            ])
+            mynp.einsum('i,jk->ik', a, b)