Explanation for weird pytree behaviour with grad?

[aniquetahir]

Consider the following code:

from collections import namedtuple
import jax 
import jax.numpy as np

Test = namedtuple('Test', ['a', 'b'])
test = Test(a=np.ones((10, 2)), b=np.ones((10,5)))

weird_function = jax.grad(lambda x: 0.)
other_function = jax.jit(lambda x: 0.)

weird_function(test)
other_function(test)

Running weird_function returns a Test instance while running other_function returns a single 0.

What is happening here, exactly? How can I make other_function also return a Test instance?

[jheek]

jax.grad returns the gradient wrt to inputs while jit simply compiles a function and returns its output. Here the argument is a Test while the output is a scalar so that explains why the types are different.