Einsum allows defining Tensors by defining their element-wise computation.
This computation is defined byequation, a shorthand form based on Einstein
summation. As an example, consider multiplying two matrices A and B to form a
matrix C. The elements of C are given by:
\[ C_{i,k} = \sum_j A_{i,j} B_{j,k} \]
or
C[i,k] = sum_j A[i,j] * B[j,k]
The corresponding einsumequationis:
ij,jk->ik
In general, to convert the element-wise equation into theequationstring,
use the following procedure (intermediate strings for matrix multiplication
example provided in parentheses):
remove variable names, brackets, and commas, (ik = sum_j ij * jk)
replace "*" with ",", (ik = sum_j ij , jk)
drop summation signs, and (ik = ij, jk)
move the output to the right, while replacing "=" with "->". (ij,jk->ik)
Many common operations can be expressed in this way. For example:
>>> m = tf.ones(2,3)
>>> e = tf.einsum('ij->ji', m0) # output[j,i] = m0[i,j]
>>> print(e.shape)
(3, 2)
Diag
>>> m = tf.reshape(tf.range(9), [3,3])
>>> diag = tf.einsum('ii->i', m)
>>> print(diag.shape)
(3,)
Trace
>>> # Repeated indices are summed.
>>> trace = tf.einsum('ii', m) # output[j,i] = trace(m) = sum_i m[i, i]
>>> assert trace == sum(diag)
>>> print(trace.shape)
()
Batch matrix multiplication
>>> s = tf.random.normal(shape=[7,5,3])
>>> t = tf.random.normal(shape=[7,3,2])
>>> e = tf.einsum('bij,bjk->bik', s, t)
>>> # output[a,i,k] = sum_j s[a,i,j] * t[a, j, k]
>>> print(e.shape)
(7, 5, 2)
This method does not support broadcasting on named-axes. All axes with
matching labels should have the same length. If you have length-1 axes,
usetf.squeezeortf.reshapeto eliminate them.
To write code that is agnostic to the number of indices in the input
use an ellipsis. The ellipsis is a placeholder for "whatever other indices
fit here".
For example, to perform a NumPy-style broadcasting-batch-matrix multiplication
where the matrix multiply acts on the last two axes of the input, use:
>>> s = tf.random.normal(shape=[11, 7, 5, 3])
>>> t = tf.random.normal(shape=[11, 7, 3, 2])
>>> e = tf.einsum('...ij,...jk->...ik', s, t)
>>> print(e.shape)
(11, 7, 5, 2)
Einsumwillbroadcast over axes covered by the ellipsis.
>>> s = tf.random.normal(shape=[11, 1, 5, 3])
>>> t = tf.random.normal(shape=[1, 7, 3, 2])
>>> e = tf.einsum('...ij,...jk->...ik', s, t)
>>> print(e.shape)
(11, 7, 5, 2)
Args
equation
astrdescribing the contraction, in the same format asnumpy.einsum.
*inputs
the inputs to contract (each one aTensor), whose shapes should
be consistent withequation.
**kwargs
- optimize: Optimization strategy to use to find contraction path using
opt_einsum. Must be 'greedy', 'optimal', 'branch-2', 'branch-all' or
'auto'. (optional, default: 'greedy').
- name: A name for the operation (optional).
Returns
The contractedTensor, with shape determined byequation.
Raises
ValueError
If
- the format ofequationis incorrect,
- number of inputs or their shapes are inconsistent withequation.
[[["Easy to understand","easyToUnderstand","thumb-up"],["Solved my problem","solvedMyProblem","thumb-up"],["Other","otherUp","thumb-up"]],[["Missing the information I need","missingTheInformationINeed","thumb-down"],["Too complicated / too many steps","tooComplicatedTooManySteps","thumb-down"],["Out of date","outOfDate","thumb-down"],["Samples / code issue","samplesCodeIssue","thumb-down"],["Other","otherDown","thumb-down"]],["Last updated 2025-08-07 UTC."],[],[],null,["# meridian.backend.einsum\n\n\u003cbr /\u003e\n\nTensor contraction over specified indices and outer product. \n\n meridian.backend.einsum(\n equation, *inputs, **kwargs\n )\n\nEinsum allows defining Tensors by defining their element-wise computation.\nThis computation is defined by `equation`, a shorthand form based on Einstein\nsummation. As an example, consider multiplying two matrices A and B to form a\nmatrix C. The elements of C are given by:\n\n\\\\\\[ C_{i,k} = \\\\sum_j A_{i,j} B_{j,k} \\\\\\]\n\nor \n\n C[i,k] = sum_j A[i,j] * B[j,k]\n\nThe corresponding einsum `equation` is: \n\n ij,jk-\u003eik\n\nIn general, to convert the element-wise equation into the `equation` string,\nuse the following procedure (intermediate strings for matrix multiplication\nexample provided in parentheses):\n\n1. remove variable names, brackets, and commas, (`ik = sum_j ij * jk`)\n2. replace \"\\*\" with \",\", (`ik = sum_j ij , jk`)\n3. drop summation signs, and (`ik = ij, jk`)\n4. move the output to the right, while replacing \"=\" with \"-\\\u003e\". (`ij,jk-\u003eik`)\n\n| **Note:** If the output indices are not specified repeated indices are summed. So `ij,jk-\u003eik` can be simplified to `ij,jk`.\n\nMany common operations can be expressed in this way. For example:\n\n**Matrix multiplication** \n\n \u003e\u003e\u003e m0 = tf.random.normal(shape=[2, 3])\n \u003e\u003e\u003e m1 = tf.random.normal(shape=[3, 5])\n \u003e\u003e\u003e e = tf.einsum('ij,jk-\u003eik', m0, m1)\n \u003e\u003e\u003e # output[i,k] = sum_j m0[i,j] * m1[j, k]\n \u003e\u003e\u003e print(e.shape)\n (2, 5)\n\nRepeated indices are summed if the output indices are not specified. \n\n \u003e\u003e\u003e e = tf.einsum('ij,jk', m0, m1) # output[i,k] = sum_j m0[i,j] * m1[j, k]\n \u003e\u003e\u003e print(e.shape)\n (2, 5)\n\n**Dot product** \n\n \u003e\u003e\u003e u = tf.random.normal(shape=[5])\n \u003e\u003e\u003e v = tf.random.normal(shape=[5])\n \u003e\u003e\u003e e = tf.einsum('i,i-\u003e', u, v) # output = sum_i u[i]*v[i]\n \u003e\u003e\u003e print(e.shape)\n ()\n\n**Outer product** \n\n \u003e\u003e\u003e u = tf.random.normal(shape=[3])\n \u003e\u003e\u003e v = tf.random.normal(shape=[5])\n \u003e\u003e\u003e e = tf.einsum('i,j-\u003eij', u, v) # output[i,j] = u[i]*v[j]\n \u003e\u003e\u003e print(e.shape)\n (3, 5)\n\n**Transpose** \n\n \u003e\u003e\u003e m = tf.ones(2,3)\n \u003e\u003e\u003e e = tf.einsum('ij-\u003eji', m0) # output[j,i] = m0[i,j]\n \u003e\u003e\u003e print(e.shape)\n (3, 2)\n\n**Diag** \n\n \u003e\u003e\u003e m = tf.reshape(tf.range(9), [3,3])\n \u003e\u003e\u003e diag = tf.einsum('ii-\u003ei', m)\n \u003e\u003e\u003e print(diag.shape)\n (3,)\n\n**Trace** \n\n \u003e\u003e\u003e # Repeated indices are summed.\n \u003e\u003e\u003e trace = tf.einsum('ii', m) # output[j,i] = trace(m) = sum_i m[i, i]\n \u003e\u003e\u003e assert trace == sum(diag)\n \u003e\u003e\u003e print(trace.shape)\n ()\n\n**Batch matrix multiplication** \n\n \u003e\u003e\u003e s = tf.random.normal(shape=[7,5,3])\n \u003e\u003e\u003e t = tf.random.normal(shape=[7,3,2])\n \u003e\u003e\u003e e = tf.einsum('bij,bjk-\u003ebik', s, t)\n \u003e\u003e\u003e # output[a,i,k] = sum_j s[a,i,j] * t[a, j, k]\n \u003e\u003e\u003e print(e.shape)\n (7, 5, 2)\n\nThis method does not support broadcasting on named-axes. All axes with\nmatching labels should have the same length. If you have length-1 axes,\nuse `tf.squeeze` or `tf.reshape` to eliminate them.\n\nTo write code that is agnostic to the number of indices in the input\nuse an ellipsis. The ellipsis is a placeholder for \"whatever other indices\nfit here\".\n\nFor example, to perform a NumPy-style broadcasting-batch-matrix multiplication\nwhere the matrix multiply acts on the last two axes of the input, use: \n\n \u003e\u003e\u003e s = tf.random.normal(shape=[11, 7, 5, 3])\n \u003e\u003e\u003e t = tf.random.normal(shape=[11, 7, 3, 2])\n \u003e\u003e\u003e e = tf.einsum('...ij,...jk-\u003e...ik', s, t)\n \u003e\u003e\u003e print(e.shape)\n (11, 7, 5, 2)\n\nEinsum **will** broadcast over axes covered by the ellipsis. \n\n \u003e\u003e\u003e s = tf.random.normal(shape=[11, 1, 5, 3])\n \u003e\u003e\u003e t = tf.random.normal(shape=[1, 7, 3, 2])\n \u003e\u003e\u003e e = tf.einsum('...ij,...jk-\u003e...ik', s, t)\n \u003e\u003e\u003e print(e.shape)\n (11, 7, 5, 2)\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ---- ||\n|------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| `equation` | a `str` describing the contraction, in the same format as `numpy.einsum`. |\n| `*inputs` | the inputs to contract (each one a `Tensor`), whose shapes should be consistent with `equation`. |\n| `**kwargs` | - optimize: Optimization strategy to use to find contraction path using opt_einsum. Must be 'greedy', 'optimal', 'branch-2', 'branch-all' or 'auto'. (optional, default: 'greedy'). - name: A name for the operation (optional). |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Returns ------- ||\n|---|---|\n| The contracted `Tensor`, with shape determined by `equation`. ||\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Raises ------ ||\n|--------------|------------------------------------------------------------------------------------------------------------------|\n| `ValueError` | If - the format of `equation` is incorrect, - number of inputs or their shapes are inconsistent with `equation`. |\n\n\u003cbr /\u003e"]]
