# Broadcasting in TensorOperations.jl (and applying ForwardDiff)

**URL:** <https://discourse.julialang.org/t/broadcasting-in-tensoroperations-jl-and-applying-forwarddiff/53421>\
**Category:** General Usage\
**Tags:** forwarddiff, tensors\
**Created:** [January 16, 2021, 1:41am UTC](https://discourse.julialang.org/t/broadcasting-in-tensoroperations-jl-and-applying-forwarddiff/53421 "2021-01-16T01:41:44Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![bhaveshshrimali](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bhaveshshrimali/32/13233_2.png) [@bhaveshshrimali](https://discourse.julialang.org/u/bhaveshshrimali)\
**Post date:** [January 16, 2021, 1:41am UTC](https://discourse.julialang.org/t/broadcasting-in-tensoroperations-jl-and-applying-forwarddiff/53421/1 "2021-01-16T01:41:44Z")

</div>

Hi everyone,

### [Edit] Question 1:

I had a quick question regarding the usage of `TensorOperations.jl`. How would one go about broadcasting certain operations such as

```julia
using TensorOperations;
function g(F)
    @tensor F[i,j] * F[i,j];
end

```

which works if I have `F` as two dimensional (say of shape `(3,3)`). What if I want to broadcast this over a `n`-dimensional array, say of shape `(3,3,100,4,...)`? In `NumPy` I would do something like:

```python3
from numpy import einsum
einsum("ij...,ij...", F, F)

```

where the ellipsis (`...`) controls the broadcasting.

I have checked `Tullio.jl` too, but I couldn’t see anything. Any pointers would be helpful.

Thanks

### Edit:

I did find `EllipsisNotation.jl` allows one to index into an array `F[1,1,..]`, but not directly of use in this case.

### Question 2:

The end goal is to be able to differentiate through `g` using `ForwardDiff` and broadcast the result. This brings me to my next question. Say I have this function:

```julia
f(u) = @. 1. + u[1]^2 + u[2]^2

```

I would take it’s gradient as

```julia
g = x-> gradient(f, x)
v = rand(2)
@assert g(v) == 2. .* v #true

```

Now what if my input to `f` is higher dimensional say of shape (2,100,4). How would I take broadcast it’s gradient over the trailing axes? I have taken a look at the posts in discourse, but if you feel this has been answered, feel free to point

Many thanks!

---

<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 16, 2021, 12:24pm UTC](https://discourse.julialang.org/t/broadcasting-in-tensoroperations-jl-and-applying-forwarddiff/53421/2 "2021-01-16T12:24:39Z")

</div>

Nothing really auto-broadcasts over extra dimensions, Julia seldom thinks of N-dimensional arrays as being collections of lower-dimensional ones. But you can make slices and apply functions to those, in various ways, for instance:

```julia
julia> using TensorOperations, TensorCast, ForwardDiff, LinearAlgebra

julia> g(F::AbstractMatrix) = @tensor F[i,j] * F[i,j]; # as in question, is just `dot`

julia> g(rand(2,3))
2.4423897829331964

julia> g3(F::AbstractArray{<:Any, 3}) = @tensor out[k] := F[i,j,k] * F[i,j,k]; # not Einstein, won't work

julia> g3(rand(2,3,4)) |> size
ERROR: TensorOperations.IndexError{String}("non-matching indices between left and right hand side:

julia> g3b(x::AbstractArray{<:Any, 3}) = g.(eachslice(x, dims=3));

julia> g3c(x::AbstractArray{<:Any, 3}) = @cast out[k] := g(x[:,:,k]); # the same thing

julia> x3 = rand(4,2,3);

julia> g3b(x3) == g3c(x3)
true

julia> g3c(x3) |> size
(3,)

julia> gn(x) = mapslices(g, x, dims=(1,2));

julia> gn(x3) |> size
(1, 1, 3)

julia> gn(rand(4, 2, 5, 6, 7)) |> size
(1, 1, 5, 6, 7)

julia> vec(gn(x3)) ≈ g3b(x3)
true

```

And gradients:

```julia
julia> ForwardDiff.gradient(g, rand(2,3))
2×3 Matrix{Float64}:
 1.6148 1.72041 1.46002
 0.840556 0.139716 0.254563

julia> grad3(x::AbstractArray{<:Any, 3}) = @cast out[k,i,j] := ForwardDiff.gradient(g, x[:,:,k])[i,j];

julia> grad3(x3) |> size
(3, 4, 2)

julia> ForwardDiff.jacobian(g3b, x3) |> size # not the same thing, doesn't know about sparsity
(3, 24)

julia> reshape(ForwardDiff.jacobian(g3b, x3), 3, 8, 3) # now take nonzero parts:

julia> vcat(ans[1:1,:,1], ans[2:2,:,2], ans[3:3,:,3]) ≈ reshape(grad3(x3), 3, :)
true

```

There has been discussion of trying to write something like Jax’s `vmap` for Julia, re-writing lower-dimensional operations into higher ones. But notice that both for `g` and its gradient, there isn’t an obvious higher-dimensional function to replace the one on matrices. That would be where you might hope such a transformation would lead to efficiency gains, by replacing for instance many dot products with one matrix multiplication. Without that, it’s just a loop over slices, which is what `mapslices` or TensorCast already does.

The notation of `mapslices` will let it work like EllipsisNotation if the `..` are at the end, as in `gn(x)` above. It doesn’t accept multiple arguments so there is no question about combining different array sizes like broadcasting does. (Doing something like that in TensorCast & Tullio is somewhere on the wish-list.)

---

<div class="post-metadata">

**Author:** ![bhaveshshrimali](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bhaveshshrimali/32/13233_2.png) [@bhaveshshrimali](https://discourse.julialang.org/u/bhaveshshrimali)\
**Post date:** [January 16, 2021, 5:59pm UTC](https://discourse.julialang.org/t/broadcasting-in-tensoroperations-jl-and-applying-forwarddiff/53421/3 "2021-01-16T17:59:36Z")

</div>

Thanks @mcabbott  
Indeed, Jax’s `vmap` was something that I had [tried before](https://github.com/google/jax/issues/4127#issue-684036072). But it turned out to be not as fast as needed for my application (Finite Elements).

---

<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 16, 2021, 6:44pm UTC](https://discourse.julialang.org/t/broadcasting-in-tensoroperations-jl-and-applying-forwarddiff/53421/4 "2021-01-16T18:44:04Z")

</div>

Maybe you know this, but if the dimensions you’re dealing with are in fact as small as 3x3, then reinterpreting as StaticArrays is another way to go. And since that gives you an array of arrays, you can broadcast over the container as normal.

---

<div class="post-metadata">

**Author:** ![bhaveshshrimali](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bhaveshshrimali/32/13233_2.png) [@bhaveshshrimali](https://discourse.julialang.org/u/bhaveshshrimali)\
**Post date:** [January 16, 2021, 9:38pm UTC](https://discourse.julialang.org/t/broadcasting-in-tensoroperations-jl-and-applying-forwarddiff/53421/5 "2021-01-16T21:38:22Z")

</div>

Thanks. I always forget about [StaticArrays](https://discourse.julialang.org/t/computing-inverse-of-a-stack-of-matrices/45580/2):). Indeed, for my application I only need 3x3 and 2x2.

This should do it for me!

Thanks again