Zygote 不支持突变数组

问：

我正在尝试使用科学机器学习框架从 Julia 中的示例中编写螺旋示例

但是，当代码通过函数进行区分时，我收到错误：get_batch

ERROR: Mutating arrays is not supported -- called setindex!(Vector{Int64}, ...) This error occurs when you ask Zygote to differentiate operations that change the elements of arrays in place (e.g. setting values with x .= ...)

下面是一个最小的可重现示例：

  using Lux, DiffEqFlux, DifferentialEquations, ComponentArrays, Random, StatsBase, MLUtils
  using Zygote

  true_y0 = [2., 0.];
  true_A = [-0.1 2.; -2. -0.1];

  data_size = 1000;
  times = LinRange(0, 25, data_size);

  function ground_truth!(du, u, p, t)
      du .= true_A * (u.^3) 
  end

  ground_truth_odeProb = ODEProblem(ground_truth!, true_y0, (0, times[end]));

  sol_ode = Array(
                    solve(ground_truth_odeProb,
                          Tsit5(),
                          abstol = 1e-10, reltol = 1e-10,
                          saveat = times)
                  );

  batch_time = 10
  batch_size = 20

  function get_batch()

    s = sample(range(1, data_size - batch_time), batch_size, replace=false)

    batch_y0 = sol_ode[ :, s];
    
    batch_t = times[1:batch_time];

    batch_y = stack([sol_ode[:, s .+ i] for i in range(1, batch_time)], dims=3);

    return [batch_y0, batch_t, batch_y];

  end

  const neural_net = Lux.Chain(Lux.Dense(2, 50, tanh),
                        Lux.Dense(50, 2))     

  rng = Random.default_rng();
  p, st = Lux.setup(rng, neural_net)

  const _st = st;

  p_init = ComponentArray(p);

  function neural_net_func!(du, u, p, t)
    du .= neural_net(u.^3, p, st)[1];
  end

  prob_nn = ODEProblem(neural_net_func!, [0., 0.], [0., 0.], p);

  function predict(θ, y0s, ts)
    
    _prob = remake( prob_nn, 
                    u0 = y0s, 
                    tspan = (ts[1], ts[end]), 
                    p = θ
                  );
    
    Array(solve(_prob, Tsit5(), saveat = ts,
              abstol = 1e-5, reltol = 1e-5)); 
  end

  function test(θ)
    y0s, ts, targets = get_batch();
    pred = predict(θ, y0s, ts);
    loss = sum(abs2, targets .- pred)
  end

  x, lambda = pullback((θ) -> test(θ), p_init,);

  lambda(x) # give the error