parameters: (axes, *inputs)
return_statement: return math_ops.reduce_sum(expanded_output, reduction_idx)
docstring: A generalized contraction between tensors of arbitrary dimension.
  
    Like numpy.einsum.
docstring_summary: A generalized contraction between tensors of arbitrary dimension.
docstring_tokens: A generalized contraction between tensors of arbitrary dimension .
function: def einsum(axes, *inputs):
    """
    A generalized contraction between tensors of arbitrary dimension.
  
    Like numpy.einsum.
    """
  
    match = re.match('([a-z,]+)->([a-z]+)', axes)
    assert match, \
      "Indices have incorrect format: %s" % axes
  
    inputs = list(inputs)
    idx_in = match.group(1).split(',')
    idx_out = match.group(2)
    idx_all = set(''.join(idx_in))
  
  
    assert len(idx_in) == len(inputs), \
      "Expected %d inputs but only got %d" % (len(idx_in), len(inputs))
  
    # transpose inputs so axes are in alphabetical order
    for i, (input_, axes_) in enumerate(zip(inputs, idx_in)):
      assert input_.get_shape().ndims == len(axes_), \
        "Input %d with axes %s has incorrect" \
        " number of dimensions (expected %d, got %d)" % (
          i, axes_, len(axes_), input_.get_shape().ndims
        )
  
      sorted_idx = sorted(axes_)
  
      if list(axes_) != sorted_idx:
        permuted = [axes_.find(ax) for ax in sorted_idx]
        inputs[i] = array_ops.transpose(input_, permuted)
        idx_in[i] = sorted_idx
  
    missing_idx = set(idx_out).difference(idx_all)
    assert not missing_idx, \
      "Unknown ouput axes: %s" % missing_idx
  
    reduction_idx = []
    shapes = [[dim if dim else -1
               for dim in tensor.get_shape().as_list()]
              for tensor in inputs]
  
    # validate shapes for broadcasting
    for j, ax in enumerate(sorted(idx_all)):
      dims = []
      for i, idx in enumerate(idx_in):
        if ax not in idx:
          shapes[i].insert(j, 1)
        else:
          dim = shapes[i][j]
          if isinstance(dim, int) and dim > 1:
            dims.append(dim)
  
      assert len(set(dims)) <= 1, \
        "Dimension mismatch on axis: %s" % ax
  
      if ax not in idx_out:
        reduction_idx.append(j)
  
    # reshape, multiply
    expanded_inputs = [array_ops.reshape(input_, shape)
                       for input_, shape in zip(inputs, shapes)]
    expanded_output = 1
    for input_ in expanded_inputs:
      expanded_output *= input_
  
    # contract
    return math_ops.reduce_sum(expanded_output, reduction_idx)
function_tokens: def einsum ( axes , * inputs ) : match = re . match ( '([a-z,]+)->([a-z]+)' , axes ) assert match , "Indices have incorrect format: %s" % axes inputs = list ( inputs ) idx_in = match . group ( 1 ) . split ( ',' ) idx_out = match . group ( 2 ) idx_all = set ( '' . join ( idx_in ) ) assert len ( idx_in ) == len ( inputs ) , "Expected %d inputs but only got %d" % ( len ( idx_in ) , len ( inputs ) ) # transpose inputs so axes are in alphabetical order for i , ( input_ , axes_ ) in enumerate ( zip ( inputs , idx_in ) ) : assert input_ . get_shape ( ) . ndims == len ( axes_ ) , "Input %d with axes %s has incorrect" " number of dimensions (expected %d, got %d)" % ( i , axes_ , len ( axes_ ) , input_ . get_shape ( ) . ndims ) sorted_idx = sorted ( axes_ ) if list ( axes_ ) != sorted_idx : permuted = [ axes_ . find ( ax ) for ax in sorted_idx ] inputs [ i ] = array_ops . transpose ( input_ , permuted ) idx_in [ i ] = sorted_idx missing_idx = set ( idx_out ) . difference ( idx_all ) assert not missing_idx , "Unknown ouput axes: %s" % missing_idx reduction_idx = [ ] shapes = [ [ dim if dim else - 1 for dim in tensor . get_shape ( ) . as_list ( ) ] for tensor in inputs ] # validate shapes for broadcasting for j , ax in enumerate ( sorted ( idx_all ) ) : dims = [ ] for i , idx in enumerate ( idx_in ) : if ax not in idx : shapes [ i ] . insert ( j , 1 ) else : dim = shapes [ i ] [ j ] if isinstance ( dim , int ) and dim > 1 : dims . append ( dim ) assert len ( set ( dims ) ) <= 1 , "Dimension mismatch on axis: %s" % ax if ax not in idx_out : reduction_idx . append ( j ) # reshape, multiply expanded_inputs = [ array_ops . reshape ( input_ , shape ) for input_ , shape in zip ( inputs , shapes ) ] expanded_output = 1 for input_ in expanded_inputs : expanded_output *= input_ # contract return math_ops . reduce_sum ( expanded_output , reduction_idx )