# Stop gradients from propagating through some branches of computation tree in Flux
**URL:** https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392
**Category:** General Usage
**Tags:** flux, machine-learning
**Created:** [January 15, 2020, 12:34pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392 "2020-01-15T12:34:07Z")
**Posts on this page:** 10
**Page:** 1
### Author:  [@e3c6](https://discourse.julialang.org/u/e3c6)
#### Post date: [January 15, 2020, 12:34pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/1 "2020-01-15T12:34:07Z")
Is it possible to perform some computation in my loss function that is excluded from the gradient, while still using the simple `Flux.train!` API?
Tensorflow and Keras have something called `stop_gradient` , which can be applied to an output to signal that it should be treated as a constant ([tf.stop\_gradient | TensorFlow v2.9.1](https://www.tensorflow.org/api_docs/python/tf/stop_gradient)). This is very handy when programming things like contrastive divergence or expectation maximization, where there is a part of the computation graph that should not be taken into account in the loss gradients.
Otherwise I can do the training loop myself, but having a `stop_gradient` in Keras was so handy that I think Flux could have something simlar?
---
### Author:  [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)
#### Post date: [January 15, 2020, 2:10pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/2 "2020-01-15T14:10:31Z")
In Flux 0.10, you can use [`Zygote.dropgrad`](https://fluxml.ai/Zygote.jl/latest/utils/#Zygote.dropgrad) for this.
---
### Author:  [@e3c6](https://discourse.julialang.org/u/e3c6)
#### Post date: [January 15, 2020, 3:33pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/3 "2020-01-15T15:33:12Z")
Thanks, this looks like just what I need.
---
### Author:  [@e3c6](https://discourse.julialang.org/u/e3c6)
#### Post date: [January 15, 2020, 3:58pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/4 "2020-01-15T15:58:12Z")
However I was expecting to be able to mutate arrays within a `dropgrad` block. Unfortunately something like this:
```julia
A = randn(2,2)
function f(x)
Zygote.dropgrad(A .= x)
return sum(A + x)
end
f'(randn(2,2))
```
doesn’t work.
---
### Author:  [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)
#### Post date: [January 15, 2020, 4:07pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/5 "2020-01-15T16:07:45Z")
What did you expect this to do? The code will still execute as normal, just no gradient will be taken. You need to put `dropgrad` on the line below, where it is actually used.
---
### Author:  [@e3c6](https://discourse.julialang.org/u/e3c6)
#### Post date: [January 15, 2020, 4:12pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/6 "2020-01-15T16:12:40Z")
You mean like this?
```julia
julia> function f(x)
A .= x
return sum(Zygote.dropgrad(A) + x)
end
julia> f'(randn(2,2))
ERROR: Can't differentiate gc_preserve_end expression
```
It still doesn’t work. I want to treat `A` as a constant when taking the gradient of `f`, and be able to mutate the contents of `A`.
---
### Author:  [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)
#### Post date: [January 15, 2020, 4:21pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/7 "2020-01-15T16:21:09Z")
Could you provide an example of what you’re actually trying to achive? You could put all the mutation in a seperate function and surround it with `dropgrad`, but it’s difficult to say in general.
---
### Author:  [@e3c6](https://discourse.julialang.org/u/e3c6)
#### Post date: [January 15, 2020, 8:06pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/8 "2020-01-15T20:06:27Z")
I have a function `f(x)`, and an external array `A`. When `f` executes it mutates `A` and then uses the mutated array to compute its output. However when the gradient is computed I want `A` to be treated as if it were a constant array. In my example above:
```julia
function f(x)
A .= x
return sum(A + x)
end
```
I expect `f'(x)` to be an array of ones (whereas if `A` is taken into account the gradient would be an array of twos).
I am programming something like contrastive divergence. In these kind of methods, you have a mutable state that you update, but then the gradient of the loss should not consider these changes. It’s the same use-case of [tf.stop\_gradient | TensorFlow v2.9.1](https://www.tensorflow.org/api_docs/python/tf/stop_gradient).
---
### Author:  [@e3c6](https://discourse.julialang.org/u/e3c6)
#### Post date: [January 15, 2020, 8:45pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/9 "2020-01-15T20:45:00Z")
Ok I found a way to do it:
```julia
A = randn(5)
mut(x) = (A .= x)
Zygote.@nograd mut
f(x) = (mut(x); sum(x .+ A))
f'(randn(5))
# returns array of ones
```
BTW it’s very cool that this example realizes the gradient is the same for all components and returns a `FillArray` of ones instead of an ordinary `Array`.
---
### Author:  [@e3c6](https://discourse.julialang.org/u/e3c6)
#### Post date: [February 4, 2020, 6:02pm UTC](https://discourse.julialang.org/t/stop-gradients-from-propagating-through-some-branches-of-computation-tree-in-flux/33392/10 "2020-02-04T18:02:24Z")
See also: [https://github.com/FluxML/Zygote.jl/pull/465](https://github.com/FluxML/Zygote.jl/pull/465)