# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. # # 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. from __future__ import annotations import collections import itertools import re import string from typing import TYPE_CHECKING, NamedTuple import numpy as np import opt_einsum from paddle import _C_ops from ..base.data_feeder import check_type, check_variable_and_dtype from ..base.framework import in_dynamic_or_pir_mode from ..base.layer_helper import LayerHelper from .linalg import matmul, transpose from .manipulation import reshape, squeeze, unsqueeze from .math import ( multiply, sum as paddle_sum, ) if TYPE_CHECKING: from collections.abc import Sequence from paddle import Tensor __all__ = [] def parse_op_labels(labelstr: str, operand: Tensor) -> str: ''' Parse labels for an input operand. Parameters ---------- labelstr: the input label string operand: the input operand Returns ------- the input operand's full label string in which all anonymous dimensions are labeled in dots. ''' # Sanity checks for c in labelstr.replace('.', ''): assert c.isalpha(), ( f"Invalid equation: {c} is not a valid label, which should be letters." ) assert labelstr.replace('...', '', 1).find('.') == -1, ( "Invalid equation: `.` is found outside of an ellipsis." ) ndims = len(operand.shape) full_labelstr = labelstr.replace('...', '.' * (ndims - len(labelstr) + 3)) assert len(full_labelstr) == ndims, ( f"Invalid equation: the label string '{labelstr}' misses dimensions." ) return full_labelstr def parse_labels(labelstr: str, operands: Sequence[Tensor]) -> list[str]: ''' Parse label strings for all input operands. Parameters ---------- labelstr: The equation's label string operands: The input operands Returns ------- list of full label strings for all input operands ''' nop_labels = labelstr.split(',') assert len(nop_labels) == len(operands), ( f"Invalid equation: the number of operands is {len(operands)}, " f"but found {len(nop_labels)} segments in the label equation." ) return list(map(parse_op_labels, nop_labels, operands)) def validate_rhs( rhs: str, input_labels: Sequence[str], n_bcast_dims: int ) -> None: ''' Check whether the equation's right hand side is valid ''' # Sanity check. if n_bcast_dims > 0: assert '...' in rhs, ( "Invalid equation: missing ellipsis in output labels." ) rhs = rhs.replace('...', '') rhs_set = set(rhs) # Hidden assumption: available labels don't include '.' assert '.' not in input_labels # Verify that output labels all come from the set of input labels non_input_labels = rhs_set.difference(input_labels) assert not non_input_labels, ( f"Invalid equation: " f"output label {sorted(non_input_labels)} not used by any input." ) # Verify that output labels are not duplicate assert len(rhs) == len(rhs_set), ( "Invalid equation: duplicate output labels are found." ) def build_view(in_labels: str, out_labels: str) -> list[int]: ''' Build an inverse map of dimension indices. Three conditions must hold for the result to be meaningful. First, no duplicate letter labels in each label string. Second, the number of dots in dimout_labels >= that in in_labels. Third, dots are contiguous in each label string. Parameters ---------- in_labels: The dimension labels to map to out_labels: The dimension labels to map from Returns ------- The inverse map from out_labels to in_labels. The length of the inverse map equals that of out_labels. -1 is filled if there's no matching input dimension for a specific label. Examples -------- in_labels = 'ij..', out_labels = '..ji' inv_map = [2, 3, 1, 0] in_labels = 'ij..', out_labels = '..kji' inv_map = [2, 3, -1, 1, 0] ''' inv_map = [-1] * len(out_labels) # First build the broadcast dimension mapping # Find the broadcast index range in out_labels r = re.search(r'\.+', out_labels) if r is not None: start, end = r.start(), r.end() s = re.search(r'\.+', in_labels) # fill the broadcast dimension indices from right to left. if s: for ax, dim in zip( range(start, end)[::-1], range(s.start(), s.end())[::-1] ): inv_map[ax] = dim # Now work on non-broadcast dimensions it = itertools.chain(range(start), range(end, len(out_labels))) else: it = iter(range(len(out_labels))) for i in it: inv_map[i] = in_labels.find(out_labels[i]) return inv_map def build_global_view( nop_labels: Sequence[str], rhs: str | None, n_bcast_dims: int ) -> tuple[str, list[list[int]], int, list[Tensor | int]]: ''' Build the global view, which is a layout of all dimension labels plus an index table that maps from the layout to the dimensions in each operand. In the global view, the dimensions are arranged such that output ones are put on the left and contraction ones are put on the right. Parameters ---------- nop_labels: The input full label strings of all input operands rhs: The equation right hand side n_bcast_dims: The maximum number of broadcast dimensions Returns ------- A tuple of g_labels, g_view, g_nout, g_count g_labels: the layout of all labels in a string g_view: the index table g_nout: the number of output dimensions g_count: the counter array for dimension contractions ''' # Put all labels in alphabetical order concat = sorted(''.join(nop_labels).replace('.', '')) labels, count = [], [] for a, b in zip(['.', *concat], concat): if a != b: labels.append(b) count.append(1) else: count[-1] += 1 if rhs is not None: validate_rhs(rhs, labels, n_bcast_dims) g_labels_out = rhs.replace('...', '.' * n_bcast_dims) else: g_labels_out = '.' * n_bcast_dims + ''.join( l for l, c in zip(labels, count) if c == 1 ) for i in range(len(count))[::-1]: if labels[i] in g_labels_out: labels.pop(i) count.pop(i) g_labels_sum = ''.join(labels) g_labels = g_labels_out + g_labels_sum g_view = [build_view(i, g_labels) for i in nop_labels] g_nout = len(g_labels_out) g_count = count return g_labels, g_view, g_nout, g_count def build_global_shape( g_view: list[list[int]], g_labels: str, op_shapes: Sequence[list[int]] ) -> tuple[list[int], list[list[bool]]]: ''' The global shape is the shape of all dimensions rearranged and broadcasting to the global view. It's a reference data structure for einsum planning. Parameters ---------- g_view: the global view op_shapes: the shapes of the all operands Returns ------- g_shape: the global shape vector g_masks: list of shape masks for each operand. A dimension's shape mask is a boolean indicating whether its size > 1, in other words, it's not squeezable ''' view_shapes: list[list[int]] = [] g_masks: list[list[bool]] = [] for view, op_shape in zip(g_view, op_shapes): view_shapes.append([op_shape[dim] if dim > -1 else 1 for dim in view]) g_set_shape: list[set[int]] = [ set(sizes_per_ax) - {1} for sizes_per_ax in zip(*view_shapes) ] non_bcastable = [ ax for ax, sizes in enumerate(g_set_shape) if len(sizes) > 1 ] assert not non_bcastable, ( f"Invalid operands: label {g_labels[non_bcastable[0]]} " f"corresponds to non-broadcastable dimensions." ) g_shape = [sizes.pop() if len(sizes) > 0 else 1 for sizes in g_set_shape] g_masks = [ [s > 1 or s == -1 for s in view_shape] for view_shape in view_shapes ] return g_shape, g_masks def has_duplicated_labels(labels: str) -> bool: ''' Returns True if there is any duplicate label. ''' labels = labels.replace('.', '') return len(labels) > len(set(labels)) def diagonalize(labels: str, operand: Tensor) -> tuple[str, Tensor]: ''' Merges dimensions with duplicate labels. For those dimensions with duplicate labels, merge them into one dimension which represents the diagonal elements. This requires the dimensions with duplicate labels are equal sized. Examples -------- 'ijj...i' would be merged into 'ij...' ''' assert not has_duplicated_labels(labels), ( 'Duplicate labels are not supported.' ) return labels, operand def plan_reduce( plan: Plan, op: int, reduce_dims: list[int], keepdim: bool ) -> None: ''' Add reduce to the plan ''' varname = f'op{op}' f = lambda var, dims: paddle_sum(var, dims, keepdim=keepdim) step = f, [varname], varname, reduce_dims plan.add_step(step) def plan_scalar_prod(plan: Plan, op1: int, op2: int) -> None: varnames = [f'op{op1}', f'op{op2}'] f = lambda var1, var2: paddle_sum(var1) * var2 # f = lambda var1, var2: var1 * var2 step = f, varnames, varnames[1] plan.add_step(step) def plan_matmul( plan: Plan, g_view: list[list[int]], op1: int, op2: int, g_supports: list[list[bool]], g_shape: list[int], I: list[int], J1: list[int], J2: list[int], K: list[int], ) -> None: ''' plan matmul ''' # Transpose and re-shape op1 and op2 in I, J1, K and I, J2, K # Then apply matmul(x, y, transpose_x=False, transpose_y=True) var1, var2 = f'op{op1}', f'op{op2}' op1_view, op2_view = (g_view[op] for op in (op1, op2)) I1 = [idx for idx in I if op1_view[idx] >= 0] I2 = [idx for idx in I if op2_view[idx] >= 0] op1_view = np.array(op1_view) op1_dims = op1_view[I1 + J1 + K] op2_view = np.array(op2_view) op2_dims = op2_view[I2 + J2 + K] op1_mask, op2_mask = (g_supports[op] for op in (op1, op2)) op1_vshape = np.array([s if m else 1 for s, m in zip(g_shape, op1_mask)]) op2_vshape = np.array([s if m else 1 for s, m in zip(g_shape, op2_mask)]) vshape = np.maximum(op1_vshape, op2_vshape) i1, i2, j1, j2, k = map(len, (I1, I2, J1, J2, K)) if any(op1_dims != np.arange(len(op1_dims))): # print(f'perm1: {perm1}') step = transpose, [var1], var1, list(op1_dims) plan.add_step(step) if any(op2_dims != np.arange(len(op2_dims))): # print(f'perm2: {perm2}') step = transpose, [var2], var2, list(op2_dims) plan.add_step(step) # Check if conditions hold for turning the operation into a matmul if ( j1 + j2 > 0 and k > 0 and -1 not in np.concatenate((op1_vshape, op2_vshape)) ): op1_shape = [ *list(op1_vshape[I]), np.prod(op1_vshape[J1]), np.prod(op1_vshape[K]), ] op2_shape = [ *list(op2_vshape[I]), np.prod(op2_vshape[J2]), np.prod(op2_vshape[K]), ] # Merge J dims and K dims by reshaping step = reshape, [var1], var1, op1_shape plan.add_step(step) step = reshape, [var2], var2, op2_shape plan.add_step(step) # Matmul step = matmul, [var1, var2], var2, False, True plan.add_step(step) # Reshape back shape = list(vshape[I + J1 + J2]) step = reshape, [var2], var2, shape plan.add_step(step) elif j1 == j2 == k == 1: # Can still do matmul even unknown shapes are present step = matmul, [var1, var2], var2, False, True plan.add_step(step) # In the rest cases we opt for ops other than matmul else: # unsqueeze operands include J1...J2... dimensions if j2: fill = list(range(i1 + j1, i1 + j1 + j2)) step = unsqueeze, [var1], var1, fill plan.add_step(step) if j1: fill = list(range(i2, i2 + j1)) step = unsqueeze, [var2], var2, fill plan.add_step(step) # In case of no dimensions to contract, do an elementwise multiply if k == 0: # make broadcast step = multiply, [var1, var2], var2 plan.add_step(step) # Contract and no join, turn into a dot elif j1 + j2 == 0 and k == 1: step = unsqueeze, [var1], var1, [-2] plan.add_step(step) step = unsqueeze, [var2], var2, [-1] plan.add_step(step) step = matmul, [var1, var2], var2 plan.add_step(step) step = squeeze, [var2], var2, [-1, -2] plan.add_step(step) elif j1 + j2 == 0 and -1 not in np.concatenate( (op1_vshape[K], op2_vshape[K]) ): assert all(op1_vshape[K] == op2_vshape[K]) step = ( reshape, [var1], var1, [*list(op1_vshape[I]), 1, np.prod(op1_vshape[K])], ) plan.add_step(step) step = ( reshape, [var2], var2, [*list(op2_vshape[I]), 1, np.prod(op2_vshape[K])], ) plan.add_step(step) step = matmul, [var1, var2], var2, False, True plan.add_step(step) step = squeeze, [var2], var2, [-1, -2] plan.add_step(step) else: step = multiply, [var1, var2], var2 plan.add_step(step) reduce_dims = list(range(-k, 0)) plan_reduce(plan, op2, reduce_dims, keepdim=False) # Wrap up, updating auxiliary data # Updating g_mask for I and J axes for ax in I + J1 + J2: op2_mask[ax] = vshape[ax] > 1 or vshape[ax] == -1 for ax in K: op2_mask[ax] = False for ax in range(len(op2_view)): op2_view[ax] = -1 dim = 0 for ax in I + J1 + J2: op2_view[ax], dim = dim, dim + 1 g_view[op2] = list(op2_view) def plan_summation( plan: Plan, g_view: list[list[int]], op1: int, op2: int, g_supports: list[list[bool]], g_shape: list[int], g_count: list[Tensor | int], n_bcast: int, ) -> None: ''' Plan various kinds of summation ''' op1_view, op2_view = g_view[op1], g_view[op2] op1_mask, op2_mask = g_supports[op1], g_supports[op2] ndim = len(op1_view) nout = ndim - len(g_count) count = [0] * nout + g_count I, K, J1, J2 = list(range(n_bcast)), [], [], [] for ax, dim1, dim2 in zip( range(n_bcast, ndim), op1_view[n_bcast:], op2_view[n_bcast:] ): if (dim1 != -1) != (dim2 != -1): if dim1 != -1: J1.append(ax) else: J2.append(ax) elif dim1 != -1: fold = int(op1_mask[ax]) + int(op2_mask[ax]) if ax >= nout and fold == count[ax]: # Ready to fold the dimensions K.append(ax) count[ax] -= fold else: I.append(ax) count[ax] -= max(fold - 1, 0) # Update g_count g_count[:] = count[nout:] # Now it's OK to merge the K dims as the same shape holds # print(f'I: {I} J1: {J1} J2: {J2} K: {K}') plan_matmul(plan, g_view, op1, op2, g_supports, g_shape, I, J1, J2, K) def rearrange(axes: list[int]) -> tuple[list[int], list[int]]: perm, fill = [], [] for ax, dim in enumerate(axes): if dim < 0: fill.append(ax) else: perm.append(dim) # Trivial permutation returns [] if all(i == dim for i, dim in enumerate(perm)): perm = [] return perm, fill def plan_broadcast( plan: Plan, operands: Sequence[Tensor], nop_axes: list[list[int]] ) -> None: ''' Plan broadcast across ''' nop = len(operands) varnames = [f'op{i}' for i in range(nop)] for i, op_axes in zip(range(nop), nop_axes): # Re-arrange the dimensions according to the global layout perm, fill = rearrange(op_axes) var = varnames[i] if perm: step = transpose, [var], var, perm plan.add_step(step) if fill: step = unsqueeze, [var], var, fill plan.add_step(step) def f(*args): expr = ' * '.join(varnames) return eval(expr, dict(zip(varnames, args))) step = f, varnames, None plan.add_step(step) class Plan: def __init__(self): self.env = {} self.steps = [] def add_step(self, step) -> None: self.steps.append(step) def get_var(self, varname): return self.env[varname] if varname in self.env else None def set_var(self, varname, var) -> None: self.env[varname] = var def show(self) -> None: res = None for f, in_varnames, out_varname, *args in self.steps: print(repr((out_varname, f, *in_varnames, *args))) return res def execute(self): res = None for f, in_varnames, out_varname, *args in self.steps: res = f(*map(self.get_var, in_varnames), *args) if out_varname: self.set_var(out_varname, res) return res def plan_einsum( operands: Sequence[Tensor], g_view: list[list[int]], g_shape: list[int], g_supports: list[list[bool]], g_count: list[Tensor | int], n_bcast: int, ) -> Plan: ''' Plans the actual execution steps. Results ------- the execution plan ''' nop = len(operands) ndim = len(g_view[0]) nout = ndim - len(g_count) # Initialize a plan with an environment plan = Plan() op_names = [f'op{i}' for i in range(nop)] list(map(plan.set_var, op_names, operands)) # In case no dimensions to combine, do broadcast straight across if not g_count: plan_broadcast(plan, operands, g_view) return plan # Down count degenerate contraction dimensions. for view, support in zip(g_view, g_supports): # To collect the down count number, we use a type casting trick down_count = [ int((d + 1) and (not s)) for d, s in zip(view[nout:], support[nout:]) ] for i, count in enumerate(down_count): g_count[i] -= count # Reduce any dimension for which g_support is set and g_count == 1 for i, view, mask in zip(range(nop), g_view, g_supports): to_reduce = [] for dim, masked, count in zip(view[nout:], mask[nout:], g_count): to_reduce.append(dim if (masked and count == 1) else -1) reduce_dims = list(filter(lambda x: x > -1, to_reduce)) if reduce_dims: plan_reduce(plan, i, reduce_dims, keepdim=True) # Unset mask and decrease g_count for the reduced dimensions for i, d in enumerate(to_reduce): ax = i + nout mask[ax] = mask[ax] and (d == -1) g_count[i] -= 0 if d == -1 else 1 # Plan the summations over the operand sequence for i in range(nop): # plan a single step if i == 0: continue # We'd like to arrange the dimensions in the following way: # [I... J... K...] # [I... J... K...] # where # I... are aligned and not to be combined immediately # J... are not aligned and not to be combined immediately # K... are aligned and should be immediately combined # At this point the non-trivial broadcast dimensions in K are already reduced # and removed. That means all K dimensions are aligned and their sizes are not 1. # We then inspect the layout of I,J,K plus the above observation to make # specialization decisions. The current strategy is set as follows: # (1) if I... J... K... are all empty, it's multiplying a scalar # (2) if K... are empty, better use a broadcast # (3) if I... J... empty and K... not empty, a vector-vector multiply (or a dot) # (4) Elsewise, either I... or J... not empty, and K... not empty, use a general matmul # Resolve the summation kind: dot, matmul or * if not any(g_supports[i - 1]): # op1 is a one element tensor. plan_scalar_prod(plan, i - 1, i) else: plan_summation( plan, g_view, i - 1, i, g_supports, g_shape, g_count, n_bcast ) # for ax, dim in enumerate(g_view[nop-1][:nout]): # assert dim == ax assert all(not masked for masked in g_supports[nop - 1][nout:]) view = g_view[-1] if any(ax != dim for ax, dim in enumerate(view[:nout])): perm = [dim for dim in view if dim >= 0] if sorted(perm) != perm: varname = f'op{nop - 1}' step = transpose, [varname], varname, perm plan.add_step(step) dim = 0 unsqueeze_dims = [] for ax, d in enumerate(view): if d != -1: view[ax], dim = dim, dim + 1 for ax, d in enumerate(view[:nout]): if d == -1: unsqueeze_dims.append(ax) if unsqueeze_dims: varname = f'op{nop - 1}' step = unsqueeze, [varname], varname, unsqueeze_dims plan.add_step(step) squeeze_dims = [dim for dim in view[nout:] if dim != -1] if squeeze_dims: # plan_reduce(plan, nop-1, reduce_dims, keepdim=False) varname = f'op{nop - 1}' step = squeeze, [varname], varname, squeeze_dims plan.add_step(step) return plan def replace_ellipsis( left_equation: str, rhs: str, *operands: Tensor ) -> tuple[str, str, list[Tensor]]: """ we replace ... as unused variables to simplify the EinsumOp implementation. """ ellipsis_strings = None max_ndim = 0 new_operands: list[Tensor] = [] unused_variables = {chr(c) for c in range(ord('a'), ord('z'))} for equ, operand in zip(left_equation.split(','), operands): ndims = len(operand.shape) - len(equ.replace("...", "")) max_ndim = max(max_ndim, ndims) for c in equ: unused_variables.discard(c) for equ, operand in zip(left_equation.split(','), operands): if '...' in equ: start_unsqueeze_idx = equ.index('...') to_squeeze_num = max_ndim - ( len(operand.shape) - len(equ.replace("...", "")) ) operand = unsqueeze( operand, axis=[i + start_unsqueeze_idx for i in range(to_squeeze_num)], ) new_operands.append(operand) ellipsis_strings = ''.join(unused_variables.pop() for _ in range(max_ndim)) if ellipsis_strings is not None: left_equation = left_equation.replace('...', ellipsis_strings) rhs = rhs.replace('...', ellipsis_strings) return left_equation, rhs, new_operands def preprocess( equation: str, *operands: Tensor ) -> tuple[str, str, list[str], list[Tensor]]: """ check equation / raise error, default right labels generation """ equation = equation.replace(" ", "") nop = len(operands) assert nop > 0, ( f"Required at least one operand in Einsum API, but received {nop}" ) # Part the equation to left hand side and right hand side lhs, *rhs = equation.lower().split('->') assert len(rhs) < 2, "Invalid equation: multiple `->` were found." labels = parse_labels(lhs, operands) # Note, we distinguish between 'ij->' and 'ij' by setting rhs to '' and None rhs = rhs[0] if rhs else None if rhs is None: rhs = rhs_inference(lhs) assert len(lhs.split(',')) == len(operands), ( f"Invalid equation: the number of operands is {len(operands)}, " f"but found {len(lhs.split(','))} segments in the label equation." ) assert not ('...' in lhs and '...' not in rhs), ( 'Invalid equation: missing ellipsis in output labels.' ) lhs, rhs, new_operands = replace_ellipsis(lhs, rhs, *operands) return lhs, rhs, labels, new_operands class Shaped(NamedTuple): shape: list[int] def parse_fake_shape( equation: str, operands: Sequence[Tensor], labels: Sequence[str] ) -> list[Shaped]: """ this shape is just used for operands planning. may differ with the original shape. for example: ... is replaced by 1 -1 is replaced by 1 Results ------- list of shape """ origin_labels = (x.strip() for x in equation.split(',')) def fake_shape(ori_label: str, label: str, op: Tensor) -> Shaped: """ 1. ori_label is the original labels, not aligned by '....' 2. if the '...' is evaluated to empty list, there is no '.' in label """ assert len(op.shape) == len(label), ( f"length of shape and length of label must be the same, but received {len(op.shape)} != {len(label)}" ) fakes = [s for i, (l, s) in enumerate(zip(label, op.shape))] fakes = list(map(abs, fakes)) # make -1 -> 1 if '.' in ori_label: fakes.insert(ori_label.index('.'), 1) return Shaped(fakes) out = list(map(fake_shape, origin_labels, labels, operands)) return out def rhs_inference(lhs: str) -> str: def is_free(key): return cnt.get(key) == 1 and key not in ['.', ','] cnt = collections.Counter(lhs) rhs = "..." if '...' in lhs else "" rhs = rhs + "".join(filter(is_free, sorted(cnt.elements()))) return rhs def gen_equation_for_opteinsum(lhs: str, rhs: str | None) -> tuple[str, str]: """ 1. gen rhs if rhs is None 2. '...' -> 'A' """ def get_used_label(counter) -> str: used = set(counter.elements()) for c in string.ascii_lowercase: if c not in used: return c raise ValueError( "You have used all `a` - `z`, there can't find a unused char for einsum optimization" ) cnt = collections.Counter(lhs) broadcast_label = get_used_label(cnt) if rhs is None: rhs = rhs_inference(lhs) lhs = lhs.replace("...", broadcast_label) rhs = rhs.replace("...", broadcast_label) return lhs + "->" + rhs, broadcast_label def einsum_v2(equation: str, *operands: Tensor) -> Tensor: """ einsum v2 implementation. 1. Implement C++ EinsumOp. 2. V2 create the EinsumOp to calculate, so just a little verify work in python. 3. V2 use opt_einsum.contract_path to optimize the multivariable einsum. """ n_op = len(operands) lhs, rhs, labels, new_operands = preprocess(equation, *operands) if n_op <= 2: return gen_einsum_op(lhs + '->' + rhs, *new_operands) shapes = parse_fake_shape(lhs, new_operands, labels) opt_equation, broadcast_label = gen_equation_for_opteinsum(lhs, rhs) _, cons = opt_einsum.contract_path(opt_equation, *shapes, einsum_call=True) var_list = new_operands for path in cons: (a, b), _, eq, *__ = path assert a > b, ( "Assume the first var_idx is smaller than the second_idx. opt_einsum can guarantee it." ) var_s = [var_list.pop(a), var_list.pop(b)] eq = eq.replace(broadcast_label, "...") var_list.append(gen_einsum_op(eq, *var_s)) assert len(var_list) == 1, ( f"There must be one elements in list, but received {len(var_list)}." ) return var_list[0] def gen_einsum_op(equation: str, *operands: Tensor) -> Tensor: """ EinsumOp Python Interface: """ if in_dynamic_or_pir_mode(): return _C_ops.einsum(operands, equation)[0] else: assert 1 <= len(operands) <= 2, "Only support two operands in EinsumOp." for inp in operands: check_variable_and_dtype( inp, 'dtype', ['float32', 'float64'], 'einsum' ) check_type(equation, 'equation', str, 'einsum') helper = LayerHelper('einsum', **locals()) out = helper.create_variable_for_type_inference(dtype=operands[0].dtype) attrs = {} attrs['equation'] = equation caches = [ helper.create_variable_for_type_inference(dtype=operands[0].dtype) for i in range(len(operands)) ] xshape = [ helper.create_variable_for_type_inference(dtype=operands[0].dtype) for i in range(len(operands)) ] helper.append_op( type='einsum', inputs={'Operands': operands}, outputs={'Out': out, "InnerCache": caches, "XShape": xshape}, attrs=attrs, ) return out def einsum(equation: str, *operands: Tensor) -> Tensor | None: r""" einsum(equation, *operands) The current version of this API should be used in dynamic graph only mode. Einsum offers a tensor operation API which allows using the Einstein summation convention or Einstein notation. It takes as input one or multiple tensors and produces as output one tensor. Einsum is able to perform a variety of tensor operations. Following lists a few: - for single operand - trace - diagonal - transpose - sum - for double operands - dot - outer - broadcasting and elementwise multiply - matrix multiply - batched matrix multiply - for many operads - broadcasting multiply - chained matrix multiply **The summation notation** - The tensor dimensions are labeled using uncased English letters. E.g., `ijk` relates to a three dimensional tensor whose dimensions are labeled i, j, and k. - The equation is `,` separated into terms, each being a distinct input's dimension label string. - Ellipsis `...` enables broadcasting by automatically converting the unlabeled dimensions into broadcasting dimensions. - Singular labels are called free labels, duplicate are dummy labels. Dummy labeled dimensions will be reduced and removed in the output. - Output labels can be explicitly specified on the right hand side of `->` or omitted. In the latter case, the output labels will be inferred from the input labels. - Inference of output labels - Broadcasting label `...`, if present, is put on the leftmost position. - Free labels are reordered alphabetically and put after `...`. - On explicit output labels - If broadcasting is enabled, then `...` must be present. - The output labels can be an empty, an indication to output as a scalar the sum over the original output. - Non-input labels are invalid. - Duplicate labels are invalid. - For any dummy label which is present for the output, it's promoted to a free label. - For any free label which is not present for the output, it's lowered to a dummy label. - Examples - '...ij, ...jk', where i and k are free labels, j is dummy. The output label string is '...ik' - 'ij -> i', where i is a free label and j is a dummy label. - '...ij, ...jk -> ...ijk', where i, j and k are all free labels. - '...ij, ...jk -> ij', an invalid equation since `...` is not present for the output. **The summation rule** The summation procedure can be outlined as follows, although the actual steps taken may vary significantly due to implementation specific optimization. - Step 1: preparation for broadcasting, that is, transposing and unsqueezing the input operands to have each resulting dimension identically labeled across all the input operands. - Step 2: broadcasting multiply all the resulting operands from step 1. - Step 3: reducing dummy labeled dimensions. - Step 4: transposing the result tensor to match the output labels. **On trace and diagonal** The trace and diagonal are planned yet unimplemented features. Args: equation (`str`): The summation terms using the Einstein summation notation. operands (`list|Tensor`): The input tensors over which to compute the Einstein summation. The number of operands should equal the number of input terms in the equation. Returns: result (`Tensor`), the result tensor. Examples: .. code-block:: python >>> import paddle >>> paddle.seed(102) >>> x = paddle.rand([4]) >>> y = paddle.rand([5]) >>> # sum >>> print(paddle.einsum('i->', x)) Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True, 1.81225157) >>> # dot >>> print(paddle.einsum('i,i->', x, x)) Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True, 1.13530672) >>> # outer >>> print(paddle.einsum("i,j->ij", x, y)) Tensor(shape=[4, 5], dtype=float32, place=Place(cpu), stop_gradient=True, [[0.26443148, 0.05962684, 0.25360870, 0.21900642, 0.56994802], [0.20955276, 0.04725220, 0.20097610, 0.17355499, 0.45166403], [0.35836059, 0.08080698, 0.34369346, 0.29680005, 0.77240014], [0.00484230, 0.00109189, 0.00464411, 0.00401047, 0.01043695]]) >>> A = paddle.rand([2, 3, 2]) >>> B = paddle.rand([2, 2, 3]) >>> # transpose >>> print(paddle.einsum('ijk->kji', A)) Tensor(shape=[2, 3, 2], dtype=float32, place=Place(cpu), stop_gradient=True, [[[0.50882483, 0.56067896], [0.84598064, 0.36310029], [0.55289471, 0.33273944]], [[0.04836850, 0.73811269], [0.29769155, 0.28137168], [0.84636718, 0.67521429]]]) >>> # batch matrix multiplication >>> print(paddle.einsum('ijk, ikl->ijl', A,B)) Tensor(shape=[2, 3, 3], dtype=float32, place=Place(cpu), stop_gradient=True, [[[0.36321065, 0.42009076, 0.40849245], [0.74353045, 0.79189068, 0.81345987], [0.90488225, 0.79786193, 0.93451476]], [[0.12680580, 1.06945944, 0.79821426], [0.07774551, 0.55068684, 0.44512171], [0.08053084, 0.80583858, 0.56031936]]]) >>> # Ellipsis transpose >>> print(paddle.einsum('...jk->...kj', A)) Tensor(shape=[2, 2, 3], dtype=float32, place=Place(cpu), stop_gradient=True, [[[0.50882483, 0.84598064, 0.55289471], [0.04836850, 0.29769155, 0.84636718]], [[0.56067896, 0.36310029, 0.33273944], [0.73811269, 0.28137168, 0.67521429]]]) >>> # Ellipsis batch matrix multiplication >>> print(paddle.einsum('...jk, ...kl->...jl', A,B)) Tensor(shape=[2, 3, 3], dtype=float32, place=Place(cpu), stop_gradient=True, [[[0.36321065, 0.42009076, 0.40849245], [0.74353045, 0.79189068, 0.81345987], [0.90488225, 0.79786193, 0.93451476]], [[0.12680580, 1.06945944, 0.79821426], [0.07774551, 0.55068684, 0.44512171], [0.08053084, 0.80583858, 0.56031936]]]) """ import os if int(os.environ.get('FLAGS_new_einsum', "1")): return einsum_v2(equation, *operands) nop = len(operands) assert nop > 0, "At least one operand is expected." # Part the equation to left hand side and right hand side lhs, *rhs = equation.lower().replace(' ', '').split('->') assert len(rhs) < 2, "Invalid equation: multiple `->` were found." # Note, we distinguish between 'ij->' and 'ij' by setting rhs to '' and None rhs = rhs[0] if rhs else None # Parse labels for each operand and count the number of occurrences for each alphabet label nop_labels = parse_labels(lhs, operands) # Diagonalize the operands which have duplicate labels nop_labels, operands = list(zip(*map(diagonalize, nop_labels, operands))) # To handle broadcasting, we should first know how many dimensions are there # We need to use that number to generate output labels # e.g. 1 for ['ij', 'i.', '.k'] n_bcast_dims = max(s.count('.') for s in nop_labels) # Build the data structures for planning. It's helpful to think of all the operands # broadcasting together from a global view. In this view, dimensions from multiple # operands are mapped to the same position if they are labeled uniquely. Broadcasting # dimensions are mapped to adjacent positions with the right bound fixed. Subject to # each operand, the map is injective but for all operands the map is on-to. # g_labels: # The labels of the global view # g_view: # Includes a list of maps from each operand's dimensions to the global view's dimensions # which we refer to as ax or axes in the code to distinguish from operand's dims # g_shape: # The shape of the global view. The size of each dimension is what the aligned dimensions # should broadcast to # g_nout: # Number of output axes # g_supports # Booleans indicating each operand's non-trivial dimensions # g_count # Counting how many non-trivial dimensions remain for each ax g_labels, g_view, g_nout, g_count = build_global_view( nop_labels, rhs, n_bcast_dims ) g_shape, g_supports = build_global_shape( g_view, g_labels, [op.shape for op in operands] ) # Now we're ready to build up an execution plan plan = plan_einsum( operands, g_view, g_shape, g_supports, g_count, n_bcast_dims ) result = plan.execute() return result