# Flux/Zygote: Gradient with respect to inputs and implicit parameters (in 2021)

**URL:** https://discourse.julialang.org/t/flux-zygote-gradient-with-respect-to-inputs-and-implicit-parameters-in-2021/71931
**Category:** Machine Learning
**Tags:** question, flux, zygote
**Created:** [November 23, 2021, 9:13am UTC](https://discourse.julialang.org/t/flux-zygote-gradient-with-respect-to-inputs-and-implicit-parameters-in-2021/71931 "2021-11-23T09:13:46Z")
**Posts on this page:** 2
**Page:** 1

<div class="post-metadata">

### Author: ![manuelbb-upb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/manuelbb-upb/32/25222_2.png) [@manuelbb-upb](https://discourse.julialang.org/u/manuelbb-upb)
#### Post date: [November 23, 2021, 9:13am UTC](https://discourse.julialang.org/t/flux-zygote-gradient-with-respect-to-inputs-and-implicit-parameters-in-2021/71931/1 "2021-11-23T09:13:46Z")

</div>

Hello all,

I have struggled a bit to get the gradient(s) of a loss with a Zygote model both with respect to the model parameters and to the model input.  
Is it even sensible to do this in one pass?  
Consider this simple setup:

```julia
using Flux
layer = Dense(2,3)
layer_params = params( layer )
loss_fn(y_pred, y) = Flux.Losses.mse(y_pred, y)

x = rand(2) # current sample
target = ones(3)

```

Now both calls work as expected:  
Taking the gradient with respect to the model parameters

```julia
gradient( () -> loss_fn( layer(x), target ), layer_params )

```

and with respect to the inputs:

```julia
gradient( ( _x ) -> loss_fn( layer( _x ), target ), x )

```

But due to the way we take gradients with respect to implicit parameters, I could not get both in one call, e.g., **this does not work** :

```julia
gradient( ( _x ) -> loss_fn( layer( _x ), target ), x, layer_params )

```

In some old discussion (which sadly I cannot find anymore) I read that you can wrap `x` as a `Flux.Params` object. However, there only appears to be a function with signature

```julia
gradient( :: Function, :: Params )

```

so that currently (for multiple `Params`) I do

```julia
input_params = params(x)
ps = union( input_params, layer_params )
gradient( () -> loss_fn( layer( x ), target ), ps )

```

**and this works.**

However, I wonder if it is performant, especially if I loop over multiple samples `x`.  
Is there some other way to achieve what I am trying. Or is it just a dumb idea?

---

<div class="post-metadata">

### Author: ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)
#### Post date: [November 23, 2021, 12:28pm UTC](https://discourse.julialang.org/t/flux-zygote-gradient-with-respect-to-inputs-and-implicit-parameters-in-2021/71931/2 "2021-11-23T12:28:14Z")

</div>

I think that your solution is the only solution at the moment.  
I also think that the overhead would be small. In the union, you essentially create a shallow copy of IdDict and that should be pretty fast, in comparison of the price of the gradient.

You can check it out by yourself. Do few iterations where you will just take gradient with respect to parameters (no union) and then of your solutions. The preformance diff will be small.
