Grad

[Marcode10]

Is it possible to calculate the gradient of a model w.r.t input?

[jheek]

It by defualt jax.grad only computes a grad w.r.t. the first argument so you get:

def loss_fn(model, x):
  ....

model_grad = jax.grad(model, batch)

but you can also do the following:

model_grad, input_grad = jax.grad(model, batch, argnums=(0, 1))

[Marcode10]

I try to define a model and calculate its gradient by following.

class LeNET(flax.nn.Module):
    def apply(self, x):
        x = flax.nn.Conv(x, features=6, kernel_size=(3, 3))
        ....
        x = flax.nn.log_softmax(x)
        return x

_, initial_params = LeNET.init_by_shape(jax.random.PRNGKey(0),[((1, 32, 32, 1), jnp.float32)])
model = flax.nn.Model(LeNET, initial_params)

when I calculate the gradient of a model w.r.t to some random input.

x = jnp.ones((1, 32,32,1))
model_grad = jax.grad(model)(x)

I get the following error
TypeError: Gradient only defined for scalar-output functions. Output had shape: (1, 10).

Can you please direct me in this example to the way of calculating the gradient of model w.r.t input, at some input?

[jheek]

The problem here is that you have multiple outputs. A gradient is only defined for scalars. So typically you compute the gradient wrt to some loss function

[Marcode10]

Yeah, I see your point. Do you advise any gradient calculator function for vector-output functions in JAX?

[jheek]

You can use jax.jacrev or jax.jacfwd to compute the Jacobian instead