parameters: (plan, g_view, op1, op2, g_supports, g_shape, g_count,
                     n_bcast)
docstring: Plan various kinds of summation
docstring_summary: Plan various kinds of summation
docstring_tokens: Plan various kinds of summation
function: def plan_summation(plan, g_view, op1, op2, g_supports, g_shape, g_count,
                     n_bcast):
      '''
      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)
function_tokens: def plan_summation ( plan , g_view , op1 , op2 , g_supports , g_shape , g_count , n_bcast ) : 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 )