Tensor contraction is an operation that can be performed on a tensor. It is a generalization of the idea of the trace of a matrix. For example, if we have a rank-2 tensor (a matrix) and we contract along both of its axes, we take the sum of each entry in which the indices of these axes are equal. In this case we sum the entries at (0,0), (1,1), (2,2), and so on to get a rank-0 tensor (a scalar), which in this case, is the trace of the matrix.

As another example, if we have a rank-4 tensor (a 4-dimensional array), contraction will output a rank-2 tensor (a matrix). If we are contracting over the first and third axes, the axes in the result correspond to the remaining axes. The entry at (3,4) of the output is produced by summing the entries of the original at (0,3,0,4), (1,3,1,4), and so on.


Given a rank-n tensor and two distinct axes, perform tensor contraction along those axes.


You may take the tensor in any reasonable format such as an n dimensional array or list-of-lists-of-...-lists. The tensor is always guaranteed to have equal length (at least 1) along every axis. You may take the axes as either 0-indexed or 1-indexed, as well as in either direction (0 or 1 may refer to the first axis in a list-of-lists or the most nested axis).


Any reasonable format for an rank-(n-2) tensor. In particular, a rank-0 scalar cannot be nested in an array.


Using 0-indexing:

[[0, 1, 2],
 [3, 4, 5],
 [6, 7, 8]], 0, 1 -> 12 # trace of the matrix

[[[4, 8, 4],
  [1, 6, 9],
  [2, 8, 2]],
 [[8, 4, 9],
  [7, 9, 2],
  [0, 4, 4]],
 [[2, 7, 7],
  [6, 3, 0],
  [3, 8, 8]]], 0, 1 -> [14, 25, 14]

[[[4, 8, 4],
  [1, 6, 9],
  [2, 8, 2]],
 [[8, 4, 9],
  [7, 9, 2],
  [0, 4, 4]],
 [[2, 7, 7],
  [6, 3, 0],
  [3, 8, 8]]], 0, 2 -> [15, 10, 14] # different answer along different axes

[[[[8, 5],
   [7, 9]],
  [[2, 0],
   [6, 0]]],
 [[[8, 3],
   [2, 2]],
  [[1, 3],
   [9, 1]]]], 0, 2 -> [[10,  7],
                       [11,  1]]

[[[[[7]]]]], 1, 4 -> [[[7]]] # contraction always reduces rank by 2

Standard loopholes are forbidden. As this is , shortest program wins.

  • \$\begingroup\$ Can we assume that we will always receive the axes inputs in sorted order? \$\endgroup\$ Commented Feb 26 at 21:56
  • \$\begingroup\$ May I take the number of dimensions of the array as an input? \$\endgroup\$
    – Neil
    Commented Feb 26 at 23:35
  • \$\begingroup\$ @CursorCoercer yes, and you may assume ascending or descending order \$\endgroup\$
    – Tbw
    Commented Feb 27 at 3:07
  • \$\begingroup\$ @Neil no, only the two axes \$\endgroup\$
    – Tbw
    Commented Feb 27 at 3:09
  • \$\begingroup\$ What is the win condition? I assume code-golf? Could you add the appropriate tag for it. \$\endgroup\$ Commented Feb 27 at 8:33

6 Answers 6


Python + NumPy, 24 bytes

-3 bytes thanks to @Mukundan314.

lambda a,i:a.trace(0,*i)

Attempt This Online!

  • 2
    \$\begingroup\$ 24 bytes: lambda a,i:a.trace(0,*i) \$\endgroup\$ Commented Feb 27 at 4:13
  • \$\begingroup\$ Nice! I didn't even know trace had this functionality. I was making all the test cases with einsum \$\endgroup\$
    – Tbw
    Commented Feb 27 at 14:57

APL (Dyalog Unicode), 23 bytes


Follows a golf suggested by Marshall and Adam. My original answer that shows the idea better:


Matrix on the left, the two axes on the right. Works by noticing that tensor contraction on two axes given a tensor A is equivalent to dyadic transposition with rank(A)-2 in place of axes and iota everywhere else. Then we sum once by trailing axis. Transposition reduces rank by one, sum on trailing axis also reduces rank by one. Hence the resulting tensor has rank of rank(A)-2. Assumes quad IO of zero.

  • \$\begingroup\$ ⍵∊⍨⍳r←⍴⍴⍺ — also, your TIO link doesn't do anything. \$\endgroup\$
    – Adám
    Commented Feb 27 at 4:31
  • \$\begingroup\$ I believe +/⊢⍉⍨-⍨⍥≢∘⍴⌊{⍋⍋⍺∊⍨⍳≢⍴⍵} works at only 23 bytes, \$\endgroup\$
    – Adám
    Commented Feb 27 at 4:39
  • \$\begingroup\$ it should say that it takes a tensor on the right and axes on the left. Here's a link to the APLgolf \$\endgroup\$
    – Tbw
    Commented Feb 27 at 16:17
  • 1
    \$\begingroup\$ +/{(~∘⍺⍳⊢)⍳≢⍴⍵}⍉⊢ \$\endgroup\$
    – att
    Commented Feb 27 at 18:55
  • 1
    \$\begingroup\$ @att The last two lines of your test cases should have 5s instead of 4s. It still works correctly when you do so. \$\endgroup\$
    – Tbw
    Commented Feb 28 at 20:20

JavaScript (Node.js), 115 bytes


Attempt This Online!

  • \$\begingroup\$ Seems 33 bytes in jcram, though i have no idea how it works: ζⓅVⓈ⊂⍕⍈AξⒹ↓⍋⏦⑧ⓀPzyⓁε⍙⁴⍨⍔D⌊W⍰⍑⍠⍁⏧⍎. And it was 43 bytes if i did not adjust variable names... May the improve be done automatically? \$\endgroup\$
    – tsh
    Commented Feb 27 at 11:51
  • \$\begingroup\$ re: the variable name thing: that's just an artifact of this jcram version being experimental and a little lacking, sorry :-). As for how it works, a bit like gzip except with superpowers. \$\endgroup\$ Commented Feb 27 at 14:05
  • \$\begingroup\$ @KamilaSzewczyk If so, maybe you want to give your model version numbers. So we can reference to specified version of jcram to avoid any further changes breaks current submissions (if anyone had already used it on PPCG). \$\endgroup\$
    – tsh
    Commented Feb 28 at 3:48
  • \$\begingroup\$ Well. If I do end up improving JCram, it will make you use something like De Bruijn indices for arguments. And it will stop being a compressor just for "vanilla JS". So if an improvement is released, I will clearly mark it as JCram2 or something. \$\endgroup\$ Commented Feb 28 at 7:07
  • \$\begingroup\$ @KamilaSzewczyk Just come up with some random idea (and does not know if it may help, just in case): For each program in your training set, adjust its variables and try to compress it with your current compressor to see if we can archive a better score. If so, update the code. After every input get updated, we now have a new training set. And then train the model again, and maybe iterate the same process again. I believe it can somehow (?) boost the compress performance. \$\endgroup\$
    – tsh
    Commented Mar 4 at 12:32

Wolfram Language (Mathematica), 14 bytes


Axes are 1-indexed.
-5 thanks to @att
Try it online!

  • \$\begingroup\$ -5, just use the function directly :) \$\endgroup\$
    – att
    Commented Mar 9 at 8:06
  • \$\begingroup\$ @att, it's a shame (( Just recently returned to CG \$\endgroup\$
    – lesobrod
    Commented Mar 9 at 8:29

Charcoal, 103 82 81 bytes


Attempt This Online! Link is to verbose version of code. Explanation:


Get the size of the tensor.


Flatten the tensor into a list and get 1 less than the number of dimensions.


Create a temporary holding list of lists for the elements of the tensor that will be summed to make the contracted tensor. This list is still the same size as the flattened tensor for now.


For each element of the tensor, calculate whether it contributes to the contracted tensor, and if so, append that element to the relevant list element.


Sum each of the lists.


Unflatten the list back into a tensor of the original dimension.


Retrieve the first element of the first element to produce the final result.

Previous 103-byte version was a backport to TIO of a 79-byte version using my experimental multidimensional indexing branch:


Try it online! Link is to verbose version of code. Explanation:


Calculate the number of dimensions of the input tensor.


Generate a list of a zero tensor of two fewer dimensions as the initial value of the result tensor.


Loop over the multidimensional indices of all the elements of the tensor where the coordinates in the two given axes are identical.


Get the parent list of the corresponding element of the result tensor.


Get the element of the input tensor and add it to that result element.


Pretty-print the final result. (This is to show that the result is a single element when the input has two dimensions, as otherwise the output would be indistinguishable.)

Explanation of the 79-byte multidimensional indexing version:


Calculate the number of dimensions of the input tensor as before.


Generate a list of a zero initial result tensor. (This is slightly shorter because Times fully vectorises on that branch but it only vectorises over single dimensional arrays on TIO.)


Loop over the multidimensional indices of the input tensor.


Get the multidimensional index of the result element.


Update that element if the coordinates in the two given axes are identical.


Pretty-print the final result.


Pyth, 38 bytes


Try it online!

Defines a function : which takes three inputs: a nested list, first axis, and second axis. This assumes that the first axis is less than the second axis.


The function operates recursively. The two axis inputs are T and Y. There are four cases to consider:

  • In the base case T=0, Y=1, and the matrix has depth 2. In this case we simply take the trace of the matrix, which is straightforward to implement.
  • In the case that T>0 we map : over the list with axes t-1 and y-1.
  • In the case that T=0 and Y>1, we transpose the matrix and call : on the transposition with axes 1 and Y.
  • And finally in the case that T=0, Y=1, and the matrix has depth greater than 2, we map transposition over the list and call : on that with axes 0, 2.
.N                                        # define :(N, T, Y)
  ?T                                      # if T != 0:
    :RtTtYN                               #   map : over N with additional arguments T-1, Y-1
           ?tY                            # else if Y > 1:
              :CN1Y                       #   :(transpose(N), 1, Y)
                   ?sIssN                 # else if matrix depth is 2
                         s                #   sum of
                          .e   N          #   map lambda k, b over the indices, values of N
                            @bk           #     b[k]
                                          # else
                                 CMN      #   map transposition over N
                                :   T2    #   :(above, T, 2)

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.