dpnp.einsum

dpnp.einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='same_kind', optimize=False)[source]

Evaluates the Einstein summation convention on the operands.

For full documentation refer to numpy.einsum.

Parameters:

  • subscripts (str) -- Specifies the subscripts for summation as comma separated list of subscript labels. An implicit (classical Einstein summation) calculation is performed unless the explicit indicator '->' is included as well as subscript labels of the precise output form.

  • *operands (sequence of {dpnp.ndarrays, usm_ndarray}) -- These are the arrays for the operation.

  • out ({dpnp.ndarrays, usm_ndarray, None}, optional) -- If provided, the calculation is done into this array.

  • dtype ({dtype, None}, optional) -- If provided, forces the calculation to use the data type specified. Default: None.

  • order ({"C", "F", "A", "K"}, optional) -- Controls the memory layout of the output. "C" means it should be C-contiguous. "F" means it should be F-contiguous, "A" means it should be "F" if the inputs are all "F", "C" otherwise. "K" means it should be as close to the layout as the inputs as is possible, including arbitrarily permuted axes. Default: "K".

  • casting ({"no", "equiv", "safe", "same_kind", "unsafe"}, optional) --

    Controls what kind of data casting may occur. Setting this to "unsafe" is not recommended, as it can adversely affect accumulations.

    • "no" means the data types should not be cast at all.

    • "equiv" means only byte-order changes are allowed.

    • "safe" means only casts which can preserve values are allowed.

    • "same_kind" means only safe casts or casts within a kind, like float64 to float32, are allowed.

    • "unsafe" means any data conversions may be done.

    Please note that, in contrast to NumPy, the default setting here is "same_kind". This is to prevent errors that may occur when data needs to be converted to float64, but the device does not support it. In such cases, the data is instead converted to float32. Default: "same_kind".

  • optimize ({False, True, "greedy", "optimal"}, optional) -- Controls if intermediate optimization should occur. No optimization will occur if False and True will default to the "greedy" algorithm. Also accepts an explicit contraction list from the dpnp.einsum_path function. Default: False.

Returns:

out -- The calculation based on the Einstein summation convention.

Return type:

dpnp.ndarray

See also

dpnp.einsum_path

Evaluates the lowest cost contraction order for an einsum expression.

dpnp.dot

Returns the dot product of two arrays.

dpnp.inner

Returns the inner product of two arrays.

dpnp.outer

Returns the outer product of two arrays.

dpnp.tensordot

Sum products over arbitrary axes.

dpnp.linalg.multi_dot

Chained dot product.

Examples

>>> import dpnp as np
>>> a = np.arange(25).reshape(5,5)
>>> b = np.arange(5)
>>> c = np.arange(6).reshape(2,3)

Trace of a matrix:

>>> np.einsum("ii", a)
array(60)
>>> np.einsum(a, [0,0])
array(60)
>>> np.trace(a)
array(60)

Extract the diagonal (requires explicit form):

>>> np.einsum("ii->i", a)
array([ 0,  6, 12, 18, 24])
>>> np.einsum(a, [0, 0], [0])
array([ 0,  6, 12, 18, 24])
>>> np.diag(a)
array([ 0,  6, 12, 18, 24])

Sum over an axis (requires explicit form):

>>> np.einsum("ij->i", a)
array([ 10,  35,  60,  85, 110])
>>> np.einsum(a, [0, 1], [0])
array([ 10,  35,  60,  85, 110])
>>> np.sum(a, axis=1)
array([ 10,  35,  60,  85, 110])

For higher dimensional arrays summing a single axis can be done with ellipsis:

>>> np.einsum("...j->...", a)
array([ 10,  35,  60,  85, 110])
>>> np.einsum(a, [Ellipsis,1], [Ellipsis])
array([ 10,  35,  60,  85, 110])

Compute a matrix transpose, or reorder any number of axes:

>>> np.einsum("ji", c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.einsum("ij->ji", c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.einsum(c, [1, 0])
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.transpose(c)
array([[0, 3],
       [1, 4],
       [2, 5]])

Vector inner products:

>>> np.einsum("i,i", b, b)
array(30)
>>> np.einsum(b, [0], b, [0])
array(30)
>>> np.inner(b,b)
array(30)

Matrix vector multiplication:

>>> np.einsum("ij,j", a, b)
array([ 30,  80, 130, 180, 230])
>>> np.einsum(a, [0,1], b, [1])
array([ 30,  80, 130, 180, 230])
>>> np.dot(a, b)
array([ 30,  80, 130, 180, 230])
>>> np.einsum("...j,j", a, b)
array([ 30,  80, 130, 180, 230])

Broadcasting and scalar multiplication:

>>> np.einsum("..., ...", 3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.einsum(",ij", 3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.einsum(3, [Ellipsis], c, [Ellipsis])
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.multiply(3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])

Vector outer product:

>>> np.einsum("i,j", np.arange(2)+1, b)
array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])
>>> np.einsum(np.arange(2)+1, [0], b, [1])
array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])
>>> np.outer(np.arange(2)+1, b)
array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])

Tensor contraction:

>>> a = np.arange(60.).reshape(3, 4, 5)
>>> b = np.arange(24.).reshape(4, 3, 2)
>>> np.einsum("ijk,jil->kl", a, b)
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])
>>> np.einsum(a, [0, 1, 2], b, [1, 0, 3], [2, 3])
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])
>>> np.tensordot(a, b, axes=([1, 0],[0, 1]))
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

Example of ellipsis use:

>>> a = np.arange(6).reshape((3, 2))
>>> b = np.arange(12).reshape((4, 3))
>>> np.einsum("ki,jk->ij", a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])
>>> np.einsum("ki,...k->i...", a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])
>>> np.einsum("k...,jk", a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])

Chained array operations. For more complicated contractions, speed ups might be achieved by repeatedly computing a "greedy" path or computing the "optimal" path in advance and repeatedly applying it, using an einsum_path insertion. Performance improvements can be particularly significant with larger arrays:

>>> a = np.ones(64000).reshape(20, 40, 80)

Basic einsum: 119 ms ± 26 ms per loop (evaluated on 12th Gen Intel® Core™ i7 processor)

>>> %timeit np.einsum("ijk,ilm,njm,nlk,abc->",a,a,a,a,a)

Sub-optimal einsum: 32.9 ms ± 5.1 ms per loop

>>> %timeit np.einsum("ijk,ilm,njm,nlk,abc->",a,a,a,a,a, optimize="optimal")

Greedy einsum: 28.6 ms ± 4.8 ms per loop

>>> %timeit np.einsum("ijk,ilm,njm,nlk,abc->",a,a,a,a,a, optimize="greedy")

Optimal einsum: 26.9 ms ± 6.3 ms per loop

>>> path = np.einsum_path(
    "ijk,ilm,njm,nlk,abc->",a,a,a,a,a, optimize="optimal"
)[0]
>>> %timeit np.einsum("ijk,ilm,njm,nlk,abc->",a,a,a,a,a, optimize=path)