Comment by krastanov

Are you talking about "Differentiating with respect to nested lists, tuples, and dicts" from that page? The comment to which you are responding covers quite a bit more. The jax documentation specifically says "standard Python containers". Zygote.jl and other less stable Julia auto-diff libraries go far beyond the built-ins and can work with structures defined by packages never designed to be used with automatic differentiation. Of course, there are limitations, but quite a bit less severe than the one in jax (and again, I am saying this while being a big jax fan).