Vectorization over subset of indices

General Usage
1

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!

1 Like
2 
@views @. a[:,:,:,i] = b[:,:,:,i] + c[:,:,:,i]
1 Like
3

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.

4

Look at using:

1 Like
5

Thanks! Will take a look.

Do you sleep?

3 Likes
6

No

3 Likes
7

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?

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,

8

Rules are already defined for Tullio:

9

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

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,