search

JuliaMath / AbstractFFTs.jl

A Julia framework for implementing FFTs
MIT License
125 stars 32 forks source link

Converting scale array in adjoint plan application fails for subarrays #112

Open gaurav-arya opened 1 year ago

gaurav-arya commented 1 year ago

The code convert(typeof(x), scale) fails when x is a subarray, e.g. in line

https://github.com/JuliaMath/AbstractFFTs.jl/blob/1cc9ca01f9f1968face7557c372f7a90ca26f693/src/definitions.jl#L641

It would also fail if x is e.g. a Zygote.OneElement, e.g. in the context of the plan reverse rule.

Jax's implementation for reference: https://github.com/google/jax/blob/c3aa3a4c312826cdfb0d617e30f0467c18ab4564/jax/_src/lax/fft.py#L211. Not terribly useful here because there they have no problems with generic broadcasting array construction

gaurav-arya commented 1 year ago

It also fails when x is a Zygote.OneElement, causing test failures in Zygote: see https://github.com/FluxML/Zygote.jl/actions/runs/5787572971/job/15684783986 and discussion in https://github.com/FluxML/Zygote.jl/pull/1444.