# TensorFlow \`matmul\` equivalent in Julia: matrix multiplication with two given tensor dimensions

**URL:** <https://discourse.julialang.org/t/tensorflow-matmul-equivalent-in-julia-matrix-multiplication-with-two-given-tensor-dimensions/33676>\
**Category:** General Usage\
**Created:** [January 22, 2020, 9:12pm UTC](https://discourse.julialang.org/t/tensorflow-matmul-equivalent-in-julia-matrix-multiplication-with-two-given-tensor-dimensions/33676 "2020-01-22T21:12:59Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![tanhevg](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tanhevg/32/12025_2.png) [@tanhevg](https://discourse.julialang.org/u/tanhevg)\
**Post date:** [January 22, 2020, 9:12pm UTC](https://discourse.julialang.org/t/tensorflow-matmul-equivalent-in-julia-matrix-multiplication-with-two-given-tensor-dimensions/33676/1 "2020-01-22T21:12:59Z")

</div>

Does Julia have an equivalent of TensorFlow `matmul` ? I need matrix multiplication using two given (or pre-defined) dimensions of the tensors, preserving other dimensions.

Tensorflow `matmul` uses the two innermost tensor dimensions for matrix multiplication, and preserves the remaining dimensions. In pseudocode, `A_ijkmn=sum_x(B_ijkmx * C_ijkxn)`, summing across dimension `x` , which is the last dim of `B` and second last dim of `C`. So, the innermost dims must agree for multiplication, and the preserved dims must be the same between the two tensors. The actual indices do not matter.

I need this to play nicely with `CuArrays` and `Zygote` . The ultimate goal is to use it for a deep learning model.

I am aware of [TensorOperations](https://github.com/Jutho/TensorOperations.jl/), and might end up using that. The drawbacks are:

- I would not want to pull in a new dependency if there is similar functionality in packages that come with Julia Distro
- Dimensions need to be specified explicitly

---

<div class="post-metadata">

**Author:** ![PetrKryslUCSD](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/petrkryslucsd/32/215825_2.png) [@PetrKryslUCSD](https://discourse.julialang.org/u/PetrKryslUCSD)\
**Post date:** [January 22, 2020, 9:23pm UTC](https://discourse.julialang.org/t/tensorflow-matmul-equivalent-in-julia-matrix-multiplication-with-two-given-tensor-dimensions/33676/2 "2020-01-22T21:23:51Z")

</div>

`mul!(c, a, b)`

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [January 22, 2020, 9:32pm UTC](https://discourse.julialang.org/t/tensorflow-matmul-equivalent-in-julia-matrix-multiplication-with-two-given-tensor-dimensions/33676/3 "2020-01-22T21:32:53Z")

</div>

I think the right term for this is batched matrix multiplication. For every value of `i,j`, there is one `mul!` to be done.

TensorOperations.jl doesn’t allow this operation. And isn’t by default Zygote-differentiable. (Although it does now do CuArrays.)

You can use [BatchedRoutines.jl](https://github.com/Roger-luo/BatchedRoutines.jl) and its Cu friend for this. But they don’t handle derivatives. IIRC you can borrow that from [Transformers.jl](https://github.com/chengchingwen/Transformers.jl) if you dig a bit.

There is also [a PR](https://github.com/FluxML/NNlib.jl/pull/100) to add some of this to NNlib.

---

<div class="post-metadata">

**Author:** ![tanhevg](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tanhevg/32/12025_2.png) [@tanhevg](https://discourse.julialang.org/u/tanhevg)\
**Post date:** [January 22, 2020, 9:42pm UTC](https://discourse.julialang.org/t/tensorflow-matmul-equivalent-in-julia-matrix-multiplication-with-two-given-tensor-dimensions/33676/4 "2020-01-22T21:42:57Z")

</div>

I think this functionality is available in [OMEinsum](https://github.com/under-Peter/OMEinsum.jl) and (possibly without the derivs) in [TensorOperations](https://github.com/Jutho/TensorOperations.jl). Both packages need specifying the dimensions explicitly, though.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [January 22, 2020, 9:47pm UTC](https://discourse.julialang.org/t/tensorflow-matmul-equivalent-in-julia-matrix-multiplication-with-two-given-tensor-dimensions/33676/5 "2020-01-22T21:47:12Z")

</div>

Yes I just saw that this may have been be added to OMEinsum too, master? (Can’t seem to install to check.) But TensorOperations does not allow this:

```julia
julia> using TensorOperations

julia> B = rand(2,2,2,2,2); C = rand(2,2,2,2,2);

julia> @tensor A[i,j,k,m,n] := B[i,j,k,m,x] * C[i,j,k,x,n]
ERROR: TensorOperations.IndexError{String}("non-matching indices between left and right hand side: \$(Expr(:(:=), :(var\"##255\"[i, j, k, m, n]), :(var\"##253\"[i, j, k, m, x] * var\"##254\"[i, j, k, x, n])))")

```

---

<div class="post-metadata">

**Author:** ![tanhevg](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tanhevg/32/12025_2.png) [@tanhevg](https://discourse.julialang.org/u/tanhevg)\
**Post date:** [January 22, 2020, 10:14pm UTC](https://discourse.julialang.org/t/tensorflow-matmul-equivalent-in-julia-matrix-multiplication-with-two-given-tensor-dimensions/33676/6 "2020-01-22T22:14:09Z")

</div>

```julia
pkg> add OMEinsum

```

Works for me.

> TensorOperations does not allow this

Weird, this must be a recently introduced bug. Their docs claim they can do it.