ChainRules.jl

ChainRules.jl copied to clipboard

`*(::AbstractVector, ::AbstractMatrix)` pullback triggers scalar indexing on the GPU

1

When multiplying a vector by a matrix, the `rrule` triggers scalar indexing for GPU arrays. ```julia julia> using Zygote, CUDA julia> CUDA.allowscalar(false) julia> x, y = rand(Float32, 6), rand(Float32, 1,...

darsnack

Pullback for `tr` produces a CPU `Diagonal` causing downstream scalar indexing on GPUs

3

Functions like `tr(A * B)` will throw scalar indexing issues in the pullback for `*` when `A` and `B` are `CuArray`s. This is because the [pullback for `tr`](https://github.com/JuliaDiff/ChainRules.jl/blob/a9a84ba6cb8aa9ce079af9401600e7c96a8aff3a/src/rulesets/LinearAlgebra/dense.jl#L197) creates a...

darsnack

I am trying to code a function that casts a Vector{Float64} of weights and biases to ComponentArray that Lux.jl can handle. I have to do this in the optimization routine....

paul188

Make the rrule for 3-arg dot lazy

6

Currently this line https://github.com/JuliaDiff/ChainRules.jl/blob/9f1817a22404259113e230bef149a54d379a660b/src/rulesets/LinearAlgebra/dense.jl#L38 leads to a loss of structure because it constructs a dense matrix from 2 vectors. We should make this a lazy outer product (like in https://github.com/SciML/LinearSolve.jl/pull/484)...

mohdibntarek

Avoid device-to-host copy in `∇getindex!`

7

Implement kernel for accumulation in `∇getindex!` for generic index types. Fallback to old implementation if KernelAbstractions backend does not support atomics (currently this is Metal.jl). This is implemented as an...

pxl-th

Avoid device-to-host copy in `∇getindex!`

2

Can we use custom kernel with atomics for [`∇getindex!(dx::AbstractGPUArray, dy, inds...)`](https://github.com/JuliaDiff/ChainRules.jl/blob/2db48540b2fd9943c4e6db92dba1ee1e8f7f8550/src/rulesets/Base/indexing.jl#L183) instead of copying everything to CPU? This way we'd be able to avoid synchronizations and we can add such...

pxl-th

`rrule` for `map` for tuples is outdated with respect to JuliaLang/julia#42216

1

https://github.com/JuliaDiff/ChainRules.jl/blob/9dd39bdbb66819ad26d19ebbe62b76c1ede56f66/src/rulesets/Base/base.jl#L243-L273 The referenced Julia issue is resolved: JuliaLang/julia#42216

nsajko

make 3-arg dot rrule partially lazy

8

This addresses #788. I had to remove the projection to make it work otherwise I get the following error due to a missing projection method. Projecting the lazy array to...

mohdibntarek

Replace nightly tests with pre-release tests

3

The nightly CI runs have [failed persistently for 9 months](https://github.com/JuliaDiff/ChainRules.jl/actions/runs/7012621446). As a result, these are really not useful. More useful would be to test against pre-releases, and we should actually...

sethaxen

Pangoraw