numpy  2.0.0
src/multiarray/einsum.c.src File Reference
#include "Python.h"
#include "structmember.h"
#include <numpy/npy_common.h>
#include <numpy/arrayobject.h>
#include <numpy/halffloat.h>
#include <npy_pycompat.h>
#include <ctype.h>
#include "convert.h"
#include "common.h"

Defines

#define PY_SSIZE_T_CLEAN
#define NPY_NO_DEPRECATED_API   NPY_API_VERSION
#define _MULTIARRAYMODULE
#define EINSUM_USE_SSE1   0
#define EINSUM_USE_SSE2   0
#define EINSUM_IS_SSE_ALIGNED(x)   ((((npy_intp)x)&0xf) == 0)
#define NPY_EINSUM_DBG_TRACING   0
#define NPY_EINSUM_DBG_PRINT(s)
#define NPY_EINSUM_DBG_PRINT1(s, p1)
#define NPY_EINSUM_DBG_PRINT2(s, p1, p2)
#define NPY_EINSUM_DBG_PRINT3(s, p1, p2, p3)

Typedefs

typedef void(* sum_of_products_fn )(int, char **, npy_intp *, npy_intp)

Functions

static void name _sum_of_products_ noplabel (int nop, char **dataptr, npy_intp *strides, npy_intp count)
static void name
_sum_of_products_contig_ 
noplabel (int nop, char **dataptr, npy_intp *NPY_UNUSED(strides), npy_intp count)
static sum_of_products_fn get_sum_of_products_function (int nop, int type_num, npy_intp itemsize, npy_intp *fixed_strides)
static int parse_operand_subscripts (char *subscripts, int length, int ndim, int iop, char *out_labels, char *out_label_counts, int *out_min_label, int *out_max_label, int *out_num_labels)
static int parse_output_subscripts (char *subscripts, int length, int ndim_broadcast, const char *label_counts, char *out_labels)
static int get_single_op_view (PyArrayObject *op, int iop, char *labels, int ndim_output, char *output_labels, PyArrayObject **ret)
static PyArrayObjectget_combined_dims_view (PyArrayObject *op, int iop, char *labels)
static int prepare_op_axes (int ndim, int iop, char *labels, int *axes, int ndim_iter, char *iter_labels)
static int unbuffered_loop_nop1_ndim2 (NpyIter *iter)
static int unbuffered_loop_nop1_ndim3 (NpyIter *iter)
static int unbuffered_loop_nop2_ndim2 (NpyIter *iter)
static int unbuffered_loop_nop2_ndim3 (NpyIter *iter)
NPY_NO_EXPORT PyArrayObjectPyArray_EinsteinSum (char *subscripts, npy_intp nop, PyArrayObject **op_in, PyArray_Descr *dtype, NPY_ORDER order, NPY_CASTING casting, PyArrayObject *out)

Variables

static sum_of_products_fn _contig_outstride0_unary_specialization_table [NPY_NTYPES]
static sum_of_products_fn _binary_specialization_table [NPY_NTYPES][5]
static sum_of_products_fn _outstride0_specialized_table [NPY_NTYPES][4]
static sum_of_products_fn _allcontig_specialized_table [NPY_NTYPES][4]
static sum_of_products_fn _unspecialized_table [NPY_NTYPES][4]

Define Documentation

#define EINSUM_IS_SSE_ALIGNED (   x)    ((((npy_intp)x)&0xf) == 0)
#define EINSUM_USE_SSE1   0
#define EINSUM_USE_SSE2   0
TODO: Only some SSE2 for float64 is implemented.
#define NPY_EINSUM_DBG_PRINT (   s)
#define NPY_EINSUM_DBG_PRINT1 (   s,
  p1 
)
#define NPY_EINSUM_DBG_PRINT2 (   s,
  p1,
  p2 
)
#define NPY_EINSUM_DBG_PRINT3 (   s,
  p1,
  p2,
  p3 
)
#define NPY_EINSUM_DBG_TRACING   0
**** PRINTF DEBUG TRACING *********
#define NPY_NO_DEPRECATED_API   NPY_API_VERSION

Typedef Documentation

typedef void(* sum_of_products_fn)(int, char **, npy_intp *, npy_intp)
end repeat*

Function Documentation

static PyArrayObject* get_combined_dims_view ( PyArrayObject op,
int  iop,
char *  labels 
) [static]
Initialize the dimensions and strides to zero
Copy the dimensions and strides, except when collapsing
The char type may be either signed or unsigned, we need it to be signed here.
If this label says to merge axes, get the actual label
If the label is 0, it's an unlabeled broadcast dimension
Update the combined axis dimensions and strides
If the label didn't say to combine axes, increment dest i
The compressed number of dimensions
static int get_single_op_view ( PyArrayObject op,
int  iop,
char *  labels,
int  ndim_output,
char *  output_labels,
PyArrayObject **  ret 
) [static]
When there's just one operand and no reduction, we can return a view into op. This calculates the view if possible.
Initialize the dimensions and strides to zero
Match the labels in the operand with the output labels
The char type may be either signed or unsigned, we need it to be signed here.
If this label says to merge axes, get the actual label
If the label is 0, it's an unlabeled broadcast dimension
The next output label that's a broadcast dimension
Find the position for this dimension in the output
If it's not found, reduction -> can't return a view
Update the dimensions and strides of the output
If we processed all the input axes, return a view
Return success, but that we couldn't make a view
static sum_of_products_fn get_sum_of_products_function ( int  nop,
int  type_num,
npy_intp  itemsize,
npy_intp fixed_strides 
) [static]
contiguous reduction
nop of 2 has more specializations
Encode the zero/contiguous strides
Inner loop with an output stride of 0
Check for all contiguous
Contiguous loop
None of the above specializations caught it, general loops
static void bool_sum_of_products_outstride0_ noplabel ( int  nop,
char **  dataptr,
npy_intp strides,
npy_intp  count 
) [static]
System Message: ERROR/3 (<string>, line 1) Document or section may not begin with a transition.

System Message: ERROR/3 (<string>, line 1) Document may not end with a transition.
begin repeat
name = byte, short, int, long, longlong,
ubyte, ushort, uint, ulong, ulonglong, half, float, double, longdouble, cfloat, cdouble, clongdouble#
#type = npy_byte, npy_short, npy_int, npy_long, npy_longlong,
npy_ubyte, npy_ushort, npy_uint, npy_ulong, npy_ulonglong, npy_half, npy_float, npy_double, npy_longdouble, npy_cfloat, npy_cdouble, npy_clongdouble#
#temptype = npy_byte, npy_short, npy_int, npy_long, npy_longlong,
npy_ubyte, npy_ushort, npy_uint, npy_ulong, npy_ulonglong, npy_float, npy_float, npy_double, npy_longdouble, npy_float, npy_double, npy_longdouble#
to = ,,,,,
System Message: SEVERE/4 (<string>, line 15)
Unexpected section title or transition.

,,,,,

npy_float_to_half,,,, ,,#

#from = ,,,,,
System Message: SEVERE/4 (<string>, line 19)
Unexpected section title or transition.

,,,,,

npy_half_to_float,,,, ,,#

#complex = 0*5,
0*5, 0*4, 1*3#
#float32 = 0*5,
0*5, 0,1,0,0, 0*3#
#float64 = 0*5,
0*5, 0,0,1,0, 0*3#
begin repeat1
#nop = 1, 2, 3, 1000# noplabel = one, two, three, any#
&#64;nop&#64; > 3 || &#64;complex
end repeat1*
end repeat*
Do OR of ANDs for the boolean type
begin repeat
#nop = 1, 2, 3, 1000# noplabel = one, two, three, any#
This is placed before the main loop to make small counts faster
begin repeat1
#i = 6, 5, 4, 3, 2, 1, 0#
end repeat1*
Unroll the loop by 8 for fixed-size nop
If the loop was unrolled, we need to finish it off
static void name _sum_of_products_contig_ noplabel ( int  nop,
char **  dataptr,
npy_intp NPY_UNUSEDstrides,
npy_intp  count 
) [static]
functions for various &#64;nop&#64;
&#64;nop&#64; == 1
static int parse_operand_subscripts ( char *  subscripts,
int  length,
int  ndim,
int  iop,
char *  out_labels,
char *  out_label_counts,
int *  out_min_label,
int *  out_max_label,
int *  out_num_labels 
) [static]
Parses the subscripts for one operand into an output of 'ndim' labels
Process the labels from the end until the ellipsis
A label for an axis
Calculate the min and max labels
If it's the first time we see this label, count it
The end of the ellipsis
A valid ellipsis
Reduce ndim to just the dimensions left to fill at the beginning
If we stopped because of an ellipsis, start again from the beginning. The length was truncated to end at the ellipsis in this case.
A label for an axis
Calculate the min and max labels
If it's the first time we see this label, count it
Set the remaining labels to 0
Find any labels duplicated for this operand, and turn them into negative offets to the axis to merge with.
In C, the char type may be signed or unsigned, but with twos complement arithmetic the char is ok either way here, and later where it matters the char is cast to a signed char.
If this is a proper label, find any duplicates of it
Search for the next matching label
The offset from next to out_labels[idim] (negative)
Search for the next matching label
static int parse_output_subscripts ( char *  subscripts,
int  length,
int  ndim_broadcast,
const char *  label_counts,
char *  out_labels 
) [static]
Parses the subscripts for the output operand into an output that requires 'ndim_broadcast' unlabeled dimensions, returning the number of output dimensions. Returns -1 if there is an error.
Count the labels, making sure they're all unique and valid
Check if it occurs again
Check that it was used in the inputs
The number of output dimensions
Process the labels from the end until the ellipsis
A label for an axis
The end of the ellipsis
A valid ellipsis
Reduce ndim to just the dimensions left to fill at the beginning
If we stopped because of an ellipsis, start again from the beginning. The length was truncated to end at the ellipsis in this case.
A label for an axis
Set the remaining output labels to 0
static int prepare_op_axes ( int  ndim,
int  iop,
char *  labels,
int *  axes,
int  ndim_iter,
char *  iter_labels 
) [static]
If it's an unlabeled broadcast dimension, choose the next broadcast dimension from the operand.
If we used up all the operand broadcast dimensions, extend it with a "newaxis"
Otherwise map to the broadcast axis
It's a labeled dimension, find the matching one
If the op doesn't have the label, broadcast it
Otherwise use it
NPY_NO_EXPORT PyArrayObject* PyArray_EinsteinSum ( char *  subscripts,
npy_intp  nop,
PyArrayObject **  op_in,
PyArray_Descr dtype,
NPY_ORDER  order,
NPY_CASTING  casting,
PyArrayObject out 
)
This function provides summation of array elements according to
the Einstein summation convention. For example:
  • trace(a) -> einsum("ii", a)
  • transpose(a) -> einsum("ji", a)
  • multiply(a,b) -> einsum(",", a, b)
  • inner(a,b) -> einsum("i,i", a, b)
  • outer(a,b) -> einsum("i,j", a, b)
  • matvec(a,b) -> einsum("ij,j", a, b)
  • matmat(a,b) -> einsum("ij,jk", a, b)

subscripts: The string of subscripts for einstein summation. nop: The number of operands op_in: The array of operands dtype: Either NULL, or the data type to force the calculation as. order: The order for the calculation/the output axes. casting: What kind of casts should be permitted. out: Either NULL, or an array into which the output should be placed.

By default, the labels get placed in alphabetical order at the end of the output. So, if c = einsum("i,j", a, b) then c[i,j] == a[i]*b[j], but if c = einsum("j,i", a, b) then c[i,j] = a[j]*b[i].

Alternatively, you can control the output order or prevent an axis from being summed/force an axis to be summed by providing indices for the output. This allows us to turn 'trace' into 'diag', for example.

System Message: ERROR/3 (<string>, line 28) Unexpected indentation.

<blockquote>

  • diag(a) -> einsum("ii->i", a)
  • sum(a, axis=0) -> einsum("i...->", a)

</blockquote>

Subscripts at the beginning and end may be specified by putting an ellipsis "..." in the middle. For example, the function einsum("i...i", a) takes the diagonal of the first and last dimensions of the operand, and einsum("ij...,jk...->ik...") takes the matrix product using the first two indices of each operand instead of the last two.

When there is only one operand, no axes being summed, and no output parameter, this function returns a view into the operand instead of making a copy.

nop+1 (+1 is for the output) must fit in NPY_MAXARGS
Parse the subscripts string into label_counts and op_labels
Move subscripts to the start of the labels for the next op
Find the number of broadcast dimensions, which is the maximum number of labels == 0 in an op_labels array.
If there is no output signature, create one using each label that appeared once, in alphabetical order
If no output was specified, always broadcast left (like normal)
Parse the output subscript string
Parse the output subscript string
Set all the op references to NULL
Process all the input ops, combining dimensions into their diagonal where specified.
If there's just one operand and no output parameter, first try remapping the axes to the output to return a view instead of a copy.
Check whether any dimensions need to be combined
The char type may be either signed or unsigned, we need it to be signed here.
If any dimensions are combined, create a view which combines them
No combining needed
Set the output op
Set up the labels for the iterator (output + combined labels). Can just share the output_labels memory, because iter_labels is output_labels with some more labels appended.
Set up the op_axes for the iterator
Set up the op_dtypes if dtype was provided
Set the op_axes for the output
Set the iterator per-op flags
Allocate the iterator
Initialize the output to all zeros and reset the iterator
System Message: ERROR/3 (<string>, line 1) Document or section may not begin with a transition.

System Message: ERROR/3 (<string>, line 1) Document may not end with a transition.
Acceleration for some specific loop structures. Note that with axis coalescing, inputs with more dimensions can be reduced to fit into these patterns.
System Message: ERROR/3 (<string>, line 1) Document or section may not begin with a transition.

System Message: ERROR/3 (<string>, line 1) Document may not end with a transition.
Get an inner loop function, specializing it based on the strides that are fixed for the whole loop.
Finally, the main loop
If the API was needed, it may have thrown an error
static int unbuffered_loop_nop1_ndim2 ( NpyIter iter) [static]
Since the iterator wasn't tracking coordinates, the loop provided by the iterator is in Fortran-order.
static int unbuffered_loop_nop1_ndim3 ( NpyIter iter) [static]
Since the iterator wasn't tracking coordinates, the loop provided by the iterator is in Fortran-order.

References NpyIter_Deallocate().

static int unbuffered_loop_nop2_ndim2 ( NpyIter iter) [static]
Since the iterator wasn't tracking coordinates, the loop provided by the iterator is in Fortran-order.
static int unbuffered_loop_nop2_ndim3 ( NpyIter iter) [static]
Since the iterator wasn't tracking coordinates, the loop provided by the iterator is in Fortran-order.

Variable Documentation

Initial value:
 {









    {NULL, NULL, NULL, NULL},


}
End of _allcontig_specialized_table
Initial value:
 {










    {NULL, NULL, NULL, NULL, NULL},


}
End of _binary_specialization_table
Initial value:
 {




    NULL,


}
These tables need to match up with the type enum
End of _contig_outstride0_unary_specialization_table
Initial value:
 {









    {NULL, NULL, NULL, NULL},


}
End of _outstride0_specialized_table
Initial value:
 {









    {NULL, NULL, NULL, NULL},


}
End of _unnspecialized_table