einsum#

ivy.einsum(equation, *operands, out=None)[source]#

Sum the product of the elements of the input operands along dimensions specified using a notation based on the Einstein summation convention.

Parameters:
  • equation (str) – A str describing the contraction, in the same format as numpy.einsum.

  • operands (Union[Array, NativeArray]) – seq of arrays, the inputs to contract (each one an ivy.Array), whose shapes should be consistent with equation.

  • out (Optional[Array], default: None) – optional output array, for writing the result to.

Return type:

Array

Returns:

ret – The array with sums computed.

Examples

With ivy.Array input:

>>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
>>> y = ivy.einsum('ii', x)
>>> print(y)
ivy.array(12)
>>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
>>> z = ivy.einsum('ij -> j', x)
>>> print(z)
ivy.array([ 9, 12, 15])
>>> A = ivy.array([0, 1, 2])
>>> B = ivy.array([[ 0,  1,  2,  3],
...                [ 4,  5,  6,  7],
...                [ 8,  9, 10, 11]])
>>> C = ivy.einsum('i,ij->i', A, B)
>>> print(C)
ivy.array([ 0, 22, 76])
>>> A = ivy.array([[1, 1, 1],
...                [2, 2, 2],
...                [5, 5, 5]])
>>> B = ivy.array([[0, 1, 0],
...                [1, 1, 0],
...                [1, 1, 1]])
>>> C = ivy.einsum('ij,jk->ik', A, B)
>>> print(C)
ivy.array([[ 2,  3,  1],
       [ 4,  6,  2],
       [10, 15,  5]])
>>> A = ivy.arange(10)
>>> B = ivy.arange(5, 15)
>>> C = ivy.einsum('i->', A)
>>> print(C)
ivy.array(45)
>>> A = ivy.arange(10)
>>> B = ivy.arange(5, 15)
>>> C = ivy.einsum('i,i->i', A, B)
>>> print(C)
ivy.array([  0,   6,  14,  24,  36,  50,  66,  84, 104, 126])
>>> A = ivy.arange(10)
>>> B = ivy.arange(5, 15)
>>> C = ivy.einsum('i,i->', A, B) # or just use 'i,i'
>>> print(C)
ivy.array(510)
>>> A = ivy.arange(10)
>>> B = ivy.arange(5, 15)
>>> C = ivy.einsum('i,j->ij', A, B)
>>> print(C)
ivy.array([[  0,   0,   0,   0,   0,   0,   0,   0,   0,   0],
       [  5,   6,   7,   8,   9,  10,  11,  12,  13,  14],
       [ 10,  12,  14,  16,  18,  20,  22,  24,  26,  28],
       [ 15,  18,  21,  24,  27,  30,  33,  36,  39,  42],
       [ 20,  24,  28,  32,  36,  40,  44,  48,  52,  56],
       [ 25,  30,  35,  40,  45,  50,  55,  60,  65,  70],
       [ 30,  36,  42,  48,  54,  60,  66,  72,  78,  84],
       [ 35,  42,  49,  56,  63,  70,  77,  84,  91,  98],
       [ 40,  48,  56,  64,  72,  80,  88,  96, 104, 112],
       [ 45,  54,  63,  72,  81,  90,  99, 108, 117, 126]])

With a mix of ivy.Array and ivy.Container inputs:

>>> x = ivy.array([0, 1, 2])
>>> y = ivy.Container(a=ivy.array([[ 0,  1,  2,  3],
...                                [ 4,  5,  6,  7],
...                                [ 8,  9, 10, 11]]),
...                   b=ivy.array([[ 0,  1,  2],
...                                [ 4,  5,  6],
...                                [ 8,  9, 10]]))
>>> z = ivy.einsum('i,ij->i', x, y)
>>> print(z)
{
    a: ivy.array([0, 22, 76]),
    b: ivy.array([0, 15, 54])
}

With ivy.Container input:

>>> x = ivy.Container(a=ivy.array([[0, 1, 0],[1, 1, 0],[1, 1, 1]]),
...                   b=ivy.array([[0, 1, 2],[4, 5, 6],[8, 9, 10]]))
>>> y = ivy.einsum('ii', x)
>>> print(y)
{
    a: ivy.array(2),
    b: ivy.array(15)
}
Array.einsum(self, equation, *operands, out=None)[source]#

ivy.Array instance method variant of ivy.einsum. This method simply wraps the function, and so the docstring for ivy.einsum also applies to this method with minimal changes.

Parameters:
  • equation (str) – A str describing the contraction, in the same format as numpy.einsum.

  • operands (Union[Array, NativeArray]) – seq of arrays, the inputs to contract (each one an ivy.Array), whose shapes should be consistent with equation.

  • out (Optional[Array], default: None) – optional output array, for writing the result to.

Return type:

Array

Returns:

ret – The array with sums computed.

Examples

>>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
>>> y = x.einsum('ii')
>>> print(y)
ivy.array(12)
>>> x = ivy.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
>>> z = x.einsum('ij -> j')
>>> print(z)
ivy.array([ 9, 12, 15])
>>> A = ivy.array([0, 1, 2])
>>> B = ivy.array([[ 0,  1,  2,  3],
...                [ 4,  5,  6,  7],
...                [ 8,  9, 10, 11]])
>>> C = A.einsum('i,ij->i', B)
>>> print(C)
ivy.array([ 0, 22, 76])
>>> A = ivy.array([[1, 1, 1],
...                [2, 2, 2],
...                [5, 5, 5]])
>>> B = ivy.array([[0, 1, 0],
...                [1, 1, 0],
...                [1, 1, 1]])
>>> C = A.einsum('ij,jk->ik', B)
>>> print(C)
ivy.array([[ 2,  3,  1],
       [ 4,  6,  2],
       [10, 15,  5]])
>>> A = ivy.arange(10)
>>> B = A.einsum('i->')
>>> print(B)
ivy.array(45)
>>> A = ivy.arange(10)
>>> B = ivy.arange(5, 15)
>>> C = A.einsum('i,i->i', B)
>>> print(C)
ivy.array([  0,   6,  14,  24,  36,  50,  66,  84, 104, 126])
>>> A = ivy.arange(10)
>>> B = ivy.arange(5, 15)
>>> C = A.einsum('i,i->', B) # or just use 'i,i'
>>> print(C)
ivy.array(510)
Container.einsum(self, equation, key_chains=None, to_apply=True, prune_unapplied=False, map_sequences=False, *, out=None)[source]#
>>> x = ivy.Container(a=ivy.array([[0, 1, 0],[1, 1, 0],[1, 1, 1]]),
...                   b=ivy.array([[0, 1, 2],[4, 5, 6],[8, 9, 10]]))
>>> y = x.einsum('ii')
>>> print(y)
{
    a: ivy.array(2),
:rtype: :py:class:`~ivy.Container`

b: ivy.array(15)

}