Zygote: Mutating arrays is not supported (small example)

Specific Domains Machine Learning
1

A small educational example:

using Zygote
using ForwardDiff

Let’s define a simple gaussian (sampled at integer coordinates)

x = -20:1.0:20
f2(xf) = 1000.0 .* exp.(-1.25 .* (x .- xf).^2)

And a simple center of mass computation to estimate the location of the peak from the samples

function cm(xc, α)
	y = f2(xc)
	
	# The next two lines are the issue
	y = y .- 15.0;
	y[y.<0] .= 0.0;
		
	y = y.^α #just a random example, could be some other function
	i = argmax(y)
	idx = i-2:1:i+2
	return sum(y[idx] .* x[idx]) / sum(y[idx])
end

And let’s say that I want to minimize the error that I make as function of alpha

function cost_cm(α)
	r = cm.(-0.5:0.01:0.5, α) .- (-0.5:0.01:0.5)
	return maximum(r)
end

Forward diff is fine

d_cost_cm(x) = ForwardDiff.derivative(cost_cm, x)
d_cost_cm(1.0)
-0.02776240571095623

Zygote is not happy

d_cost_cm_Zygote(x) = gradient(cost_cm, x)[1]
d_cost_cm_Zygote(1.0)
Mutating arrays is not supported ...

The issue are the lines:

	y = y .- 15.0;
	y[y.<0] .= 0.0;

Is there a way to replace them with something equivalent that would make Zygote happy?
(PS do we expect this issue to be solved by Diffractor.jl?)

Somehow related, if you use filter to create a new array instead of y[y.<0] .= 0.0 can you get the list of indexes to be used in sum(y_new_filtered_array .* x[idx]) / sum(y_new_filtered_array) ?

2

You could replace y[y.<0] .= 0.0 by y = clamp.(y, 0.0, Inf).
See here for workarounds when it is more difficult to write code without mutating arrays.

3

Thank you. For fun I also created a couple of alternatives that seem to work:

function cm(xc, α)
	y0 = f2(xc)
	y0 = y0 .- 15.0;
	
	idx_m = argmax(y0)
	idx = idx_m-2:1:idx_m+2
	
	y = y0[idx];
	y = clamp.(y, 0.0, Inf)
	y = y.^α
	
	return sum(y .* x[idx]) / sum(y)
end

or

function cm(xc, α)
	y0 = f2(xc)
	y0 = y0 .- 15.0;
	
	idx = findall(x -> x>=0.0,y0)
	y = y0[idx];
	y = y.^α
	
	return sum(y .* x[idx]) / sum(y)
end
4

I suspect I have to get a bit more familiar with what “mutating” an array means here.