Reduce the verbosity of treedef printing for custom nodes.

[copybara-service]

Reduce the verbosity of treedef printing for custom nodes.

For very large trees of custom nodes this printing can be very verbose with a
lot or repetition. Our internal repository also encourages very deep package
names which exacerbates this issue.

Users encounter treedef printing when interacting with some staging APIs in JAX,
for example:

>>> params = { .. some params .. }
>>> f = jax.jit(..).lower(params).compile()
>>> f(params)  # fine
>>> params['some_new_thing'] = something
>>> f(params)
TypeError: function compiled for {treedef}, called with {treedef}.