How to use/initialize Zygote.Buffer

This is probably a common question, but when i use Zygote.Buffer i cannot assign values by doing buffer .= old_data
It throws the following:

ERROR: MethodError: no method matching ndims(::Type{Zygote.Buffer{Float64, Matrix{Float64}}})
Closest candidates are:
  ndims(::CUDA.CUSPARSE.CuSparseMatrix) at ~/.julia/packages/CUDA/DfvRa/lib/cusparse/array.jl:225
  ndims(::AbstractFFTs.Plan) at ~/.julia/packages/AbstractFFTs/Wg2Yf/src/definitions.jl:15
  ndims(::ArrayInterfaceCore.ArrayIndex{N}) where N at ~/.julia/packages/ArrayInterfaceCore/7kMjZ/src/ArrayInterfaceCore.jl:578

The relevant lines are:

    image_a = zeros(sensor.image_size)
    image = Zygote.Buffer(image_a, sensor.image_size[1], sensor.image_size[2])
    image .= image_a

When i remove the last line, it runs smoothly (at least without gradient), but obviously the image is uninitialized.

Does somebody have an idea what goes wrong here?

This is especially weird, since when i try to show ndims(image) it shows “2” and works as expected

[052768ef] CUDA v3.12.0
[31c24e10] Distributions v0.25.67
[5789e2e9] FileIO v1.15.0
[587475ba] Flux v0.13.5
[82e4d734] ImageIO v0.6.6
[6218d12a] ImageMagick v1.2.2
[a98d9a8b] Interpolations v0.14.5
[e88e6eb3] Zygote v0.6.44
[37e2e46d] LinearAlgebra
[9a3f8284] Random

Julia version 1.7.3

I solved the issue by putting two loops around the Buffer and manually assigning zero to each cell. This is a workaround, but seems to work fine so far.

Zygote.Buffer doesn’t support in-place broadcast (.=), but you can do image[:] = image_a or copyto!(image, image_a). Both will be much, much faster than setting individual elements in a loop.

However, this particular example doesn’t appear to need any of the aforementioned options? Creating a Buffer using Buffer(image_a, false) or bufferfrom(image_a) (both are equivalent) does everything in one step.

1 Like

bufferfrom is probably the function need. Thanks for the reply.