Incomplete documentation of jax.grad?

[hericks]

Hi, I am currently trying to reproduce the fundamentals of equinox in a minimal example. For this, I took another look at the documentation of jax.grad and have two questions.

1. The documentation states for the fun argument:

   fun (Callable) – Function to be differentiated. Its arguments at positions specified by argnums should be arrays, scalars, or standard Python containers. Argument arrays in the positions specified by argnums must be of inexact (i.e., floating-point or complex) type. It should return a scalar (which includes arrays with shape () but not arrays with shape (1,) etc.)

   More generally, funs arguments at positions specified by argnums can be pytrees with inexact arrays and scalars as leave nodes, correct? If yes, should this be added to the documentation? I would be able to open a quick PR.

2. The documentation states for the return value:

   A function with the same arguments as fun, that evaluates the gradient of fun. If argnums is an integer then the gradient has the same shape and type as the positional argument indicated by that integer. If argnums is a tuple of integers, the gradient is a tuple of values with the same shapes and types as the corresponding arguments. If has_aux is True then a pair of (gradient, auxiliary_data) is returned.

   Consider the following simple example:

   import jax
   from dataclasses import dataclass
   
   @dataclass
   class MyTree:
       x: float
   
   jax.tree_util.register_dataclass(MyTree, ["x"], [])
   
   def f(tree):
       return tree.x
   
   t = MyTree(x=1.0)
   print(type(t), type(t.x), t)
   
   g = jax.grad(f)(t)
   print(type(g), type(g.x), g)

   with output

   <class '__main__.MyTree'> <class 'float'> MyTree(x=1.0)
   <class '__main__.MyTree'> <class 'jaxlib.xla_extension.ArrayImpl'> MyTree(x=Array(1., dtype=float32, weak_type=True))
   

   Especially, the gradient has the same type (MyTree) as the positional argument indicated by argnums. Still, the type of the leave node x changed from float an jax.Array of dtype float32. Is this close enough or should such "arrayifications" be mentioned by the documentation?

[jakevdp]

Hi - I think you're right that the documentation is incomplete when it comes to support of general pytree arguments. This is probably because it's difficult to capture the full flexibility of the API in a docstring that's simple enough to be clear for the most common cases.

Would you like to put together a pull request to address this? If not someone from the team can take a look.