The code convert(typeof(x), scale)
fails when x
is a subarray, e.g. in line
|
return p.p \ (x ./ convert(typeof(x), scale)) |
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