# Vectorization over subset of indices

**URL:** <https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097>\
**Category:** General Usage\
**Created:** [December 25, 2022, 4:39am UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097 "2022-12-25T04:39:04Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 25, 2022, 4:39am UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097/1 "2022-12-25T04:39:05Z")

</div>

I would like to vectorize some operation such as  
a .= b .+ c  
but only a subset of indices. For example, if they are both 4D arrays, I would like to vectorize only over dimensions 1 and 3. Is this possible,perhaps with views and other methods? Thanks!

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 25, 2022, 5:23am UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097/2 "2022-12-25T05:23:38Z")

</div>

```julia
@views @. a[:,:,:,i] = b[:,:,:,i] + c[:,:,:,i]

```

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 25, 2022, 5:50am UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097/3 "2022-12-25T05:50:54Z")

</div>

Yes, but means you preallocated “a” and then you are mutatng, and Zygote will not work. But you did answer the question. I am also looking at TensorCast for ideas. I have a feeling that my problem can be solved at the cost of many allocatiins, which seems to be a consequence of non-mutation.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 25, 2022, 5:53am UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097/4 "2022-12-25T05:53:45Z")

</div>

Look at using:

> **[GitHub - mcabbott/Tullio.jl: ⅀](https://github.com/mcabbott/Tullio.jl)**
>
> ⅀. Contribute to mcabbott/Tullio.jl development by creating an account on GitHub.

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 25, 2022, 6:09am UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097/5 "2022-12-25T06:09:54Z")

</div>

Thanks! Will take a look.

Do you sleep?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 25, 2022, 6:12am UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097/6 "2022-12-25T06:12:21Z")

</div>

No

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 25, 2022, 5:45pm UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097/7 "2022-12-25T17:45:50Z")

</div>

Here is my solution, @ChrisRackauckas, using `@tullio, which looks great! I could not find where differentiation rules were defined in `Chainrules`, or perhaps they don’t have to be given only built in derivations are used?

```julia
using Zygote # use before Tulio (see docs)
using Tullio

function test(N, degree)
    x = rand(N)
    y = rand(N)

    NT = NamedTuple{(:x,:y), Tuple{Int64, Int64}}
    I = Vector{NT}()
    for i in 0:degree
        for j in 0:degree
            if i + j < (degree+1)
                push!(I, (x=i, y=j))
            end
        end
    end

    coef = rand(length(I))

    Poly(coef) = @tullio poly[i] := coef[j] * x[i]^I[j].x * y[i]^I[j].y grad=Dual nograd=x nograd=y nograd=I
    return Poly, coef
end

N = 10_000
degree =20 
@time poly, coef = test(N, degree);
@time poly(coef)
loss = x -> sum(poly(x).^2)
@time Zygote.gradient(loss, coef)
println(length(coef))

```

For `N=10_000` and degree 20 (231 terms in the polynomial), the gradient takes 0.02 sec on my MacBook Pro M1. Seems fast, but I do not have a sense of what is possible. The test functions takes 0.0002 seconds. My own problem requires degree 4 and about 200 points, which gives a Zygote timing of 0.0001.

Yes, I know I should use the benchmarking tools, shower, @time, when run 3-5 times in a row, is more than sufficient to provide an estimate, since I am interested only in orders of magnitude.

Cheers,

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 25, 2022, 5:46pm UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097/8 "2022-12-25T17:46:40Z")

</div>

Rules are already defined for Tullio:

> <https://github.com/mcabbott/Tullio.jl/blob/master/src/grad/zygote.jl>

---

<div class="post-metadata">

**Author:** ![erlebach](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/erlebach/32/12973_2.png) [@erlebach](https://discourse.julialang.org/u/erlebach)\
**Post date:** [December 25, 2022, 7:30pm UTC](https://discourse.julialang.org/t/vectorization-over-subset-of-indices/92097/9 "2022-12-25T19:30:22Z")

</div>

After all my changes, I still get the mutation error when taking derivatives with Zygote.  
More precisely,

```julia
function main(tstate::Lux.Training.TrainState, vjp::Lux.Training.AbstractVJP, data::Tuple, epochs::Int)
    #data = data .|> gpu
    # no batches? 
    for epoch in 1:epochs
        # Try computing gradient of loss
        println(data)
        grads, loss, stats, tstate = Lux.Training.compute_gradients(vjp, loss_function, data, tstate)
        @info epoch=epoch loss=loss
        tstate = Lux.Training.apply_gradients(tstate, grads)
    end
    return tstate
end

```

I get a mutation error when calling `Lux.Training.compute_gradients(tstate, grads)`.  
So the question I have is: would you expect this to fail or succeed similarly to a direct call to `Zygote.gradient(...)`? I ask before the later call works. I will spend this afternoon on this.

Cheers,