Q: how to make `f[...](jax)` default to `f[...]` ?

[nstarman]

For any object with a custom __pdoc__ is there a way to make jax (or a different array type) the default and not bracketed in parentheses?

[patrick-kidger]

Yup! You can have this method, or wadler_lindig.p{print,format,doc}(..., custom=...), return any kind of WL doc.

See for example equinox.tree_pformat, which does exactly this for JAX arrays: https://github.com/patrick-kidger/equinox/blob/d3d4eb9758a017b1caddbc5219fd09379e171681/equinox/_pretty_print.py#L46-L54

[nstarman]

Thanks! That custom looks useful, I'll probably copy it over