Indexing

[mcabbott]

This wants to add a rule for A[i,j,k] any AbstractArray.

The earlier rule was only for Array. I think the argument for that is that, ultimately, indexing of any linear algebra wrapper resolves to indexing the underlying array. But it resolves to scalar indexing, which I think will be quite inefficient for something like gradient(x -> sum(x[:,1]), transpose(rand(3,3)))[1]. And in practice that fails (with just the Array rule) as it creates & mutates an array to hold the parts.

The internal function _zerolike_writeat which was previously used for some other rules is re-named ∇getindex and simplified: I am not sure why it needed dims . It has rules to allow higher derivatives.

It also always makes a full dense array; we could consider adding something like Zygote.OneElement to be more efficient at scalar indexing. But once you add two of those you get an Array; perhaps InplaceableThunk is eventually going to be better?

[mcabbott]

This second derivative function doesn't seem to infer well, can it be improved?

[oxinabox]

Cool cool once these comments are addressed as you feel best
merge when happy

[oxinabox]

Suggested change

	"""
	_setindex_zero(x::AbstractArray, dy, inds...)
	Basically this function is like `zero(x)` except that it ensure that
	it is possible to set the value at index `inds` to `dy`.
	It does this while preserving at least the outermost the structure of `x`.
	Like `zero(x)`, it promises that `x == x + _setindex_zero(x, dy, inds...)` for all inputs; i.e. it always returns an additive identity.
	"""
	function _setindex_zero(x::AbstractArray, dy, inds...) end

I wonder if we shouldn't call it _setable_zero or something?

[mcabbott]

Maybe, I don't like any names. I guess it takes indices in a way that's like set/getindex functions.

[oxinabox]

This doesn't need any mode does it?

Suggested change

	function rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(Base.unsafe_getindex), x::AbstractVector, i::Integer)
	function rrule(cfg::RuleConfig, ::typeof(Base.unsafe_getindex), x::AbstractVector, i::Integer)

Not sure why we aren't just calling the rrule for get infact?
Or putting a Union{typeof(getindex), typeof(Base.getindex)) in the function arg on them.
(might even be able to stick view in that union too?)

[mcabbott]

I guess I thought the first Zygote use of this might leave its own rule intact for scalar indexing, to make OneElement for that, and use CR for all others. In which case rrule_via_ad here will call that.

More generally it may want to do other more efficient things for indexing ranges. Or we may want to do that here, and remove this entirely.

[oxinabox]

Can we reproduce and open an issue on Julia itself and link back here?

[mcabbott]

After fiddling a bit, here's a better version. These steps run in global scope, but fail inside the let block. FiniteDifferences is involved:

julia> let
               x_23_gpu = jl(rand(2, 3))  # using JLArrays, loaded for @gpu in test_helpers.jl
           
               # Scalar indexing, copied from:  @macroexpand @allowscalar A[i]
               y1, bk1 = rrule(CFG, Base.task_local_storage, () -> x_23_gpu[1], :ScalarIndexing, ScalarAllowed)
                @test y1 == @allowscalar x_23_gpu[1]
                bk1(1.0)
       end
ERROR: StackOverflowError:
Stacktrace:
     [1] to_vec(x::JLArray{Float64, 2})
       @ FiniteDifferences ~/.julia/packages/FiniteDifferences/VpgIT/src/to_vec.jl:73
     [2] to_vec(x::Base.ReshapedArray{Float64, 1, JLArray{Float64, 2}, Tuple{}})
       @ FiniteDifferences ~/.julia/packages/FiniteDifferences/VpgIT/src/to_vec.jl:84
--- the last 2 lines are repeated 39990 more times ---
 [79983] to_vec(x::JLArray{Float64, 2})
       @ FiniteDifferences ~/.julia/packages/FiniteDifferences/VpgIT/src/to_vec.jl:73

At best FiniteDifferences can give zero here. But the parameter which needs tracking is embedded in the function () -> x_23_gpu[1] and I don't think it can unpack that.