flax.core.copy

[matteoguarrera]

I am not quite familiar with flax and this command has been deprecated in later versions so it lacks of documentation.
How are you using this copy methods? Is it for stopping gradients?

Neural-Spectral-Methods/src/train.py

Line 18 in 090e7a1

loss = self.mod.apply(variable.copy(dict(params=params)), ϕ, method="loss")

Do you know how to make it compatible with the latest version of flax?
Thank you

[mrlazy1708]

Dear Matteo,

This line is for replacing the old params with those passed as arguments (and therefore creating gradients for jax.grad). The copy method here simply overrides the dict.

In new versions of flax, this line should be equivalent to

loss = self.mod.apply(dict(variable, params=params), ϕ, method="loss")