The example on README seems unable to work

[pollycoder]

I just tried the example on your README and now it shows the error below:

Traceback (most recent call last):
  File "/Users/polly/Desktop/SURF/surf-neural-astro/test_autofd.py", line 17, in <module>
    dFdf = jax.grad(F)(f)
           ^^^^^^^^^^^^^^
TypeError: Argument 'f[] -> f[]' of type <class 'autofd.general_array.function'> is not a valid JAX type.
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Could you please tell me how to fix this bug ? Thanks a lot !