Improve element-wise einsum
@@ -315,6 +315,18 @@ def _get_out_shape(shape0, sub0, shape1, sub1, sub_out):
return out_shape
+def _reshape_transpose(arr, sub, sub_out):
+ sub = list(sub)
+ shape = list(arr.shape)
+ axes = []
+ for i in sub_out:
+ if i not in sub:
+ sub.append(i)
+ shape.append(1)
+ axes.append(sub.index(i))
+ return cupy.transpose(arr.reshape(shape), axes)
+
+
def reduced_binary_einsum(arr0, sub0, arr1, sub1, sub_others):
set0 = set(sub0)
set1 = set(sub1)
@@ -340,6 +352,13 @@ def reduced_binary_einsum(arr0, sub0, arr1, sub1, sub_others):
sub_out = sub_b + sub_l + sub_r
assert set(sub_out) <= set_others, 'operands should be reduced: unary sum'
+ if len(contract_dims) == 0:
+ if len(sub_out) == len(sub_others):
+ sub_out = sub_others
+ arr0 = _reshape_transpose(arr0, sub0, sub_out)
+ arr1 = _reshape_transpose(arr1, sub1, sub_out)
+ return arr0 * arr1, sub_out
+
if _use_cutensor(arr0.dtype, sub0, arr1.dtype, sub1,
batch_dims, contract_dims):
if len(sub_out) == len(sub_others):
