How to use/initialize Zygote.Buffer

Specific Domains Machine Learning
1

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?

Edit:
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

2

Update:
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.

3

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.

2 Likes
4

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