Automatic Differentiation for Tensor Algebras

Kjolstad et. al. proposed a tensor algebra compiler. It takes expressions that define a tensor element-wise, such as $f{ij}(a,b,c,d) = \exp\left[-\sum{k=0}⁴ \left((a{ik}+b{jk})^2\, c{ii} + d{i+k}³ \right) \right]$, and generates the corresponding compute kernel code. For machine learning, especially deep learning, it is often necessary to compute the gradient of a loss function $l(a,b,c,d)=l(f(a,b,c,d))$ with respect to parameters $a,b,c,d$. If tensor compilers are to be applied in this field, it is necessary to derive expressions for the derivatives of element-wise defined tensors, i.e. expressions for $(da){ik}=\partial l/\partial a{ik}$. When the mapping between function indices and argument indices is not 1:1, special attention is required. For the function $f{ij} (x) = xi^2$, the derivative of the loss is $(dx)i=\partial l/\partial xi=\sumj (df){ij}2xi$; the sum is necessary because index $j$ does not appear in the indices of $f$. Another example is $f{i}(x)=x{ii}^2$, where $x$ is a matrix; here we have $(dx){ij}=\delta{ij}(df)i2x{ii}$; the Kronecker delta is necessary because the derivative is zero for off-diagonal elements. Another indexing scheme is used by $f{ij}(x)=\exp x{i+j}$; here the correct derivative is $(dx){k}=\sumi (df){i,k-i} \exp x_{k}$, where the range of the sum must be chosen appropriately. In this publication we present an algorithm that can handle any case in which the indices of an argument are an arbitrary linear combination of the indices of the function, thus all the above examples can be handled. Sums (and their ranges) and Kronecker deltas are automatically inserted into the derivatives as necessary. Additionally, the indices are transformed, if required (as in the last example). The algorithm outputs a symbolic expression that can be subsequently fed into a tensor algebra compiler. Source code is provided.