drinx.tree_diagram

drinx.tree_diagram(tree, max_depth=None, static_leaves=False)

Render a JAX pytree as an ASCII tree diagram.

Parameters:

• tree (Any) – Any JAX pytree.

• max_depth (int | None) – Maximum depth to expand. None means unlimited.

• static_leaves (bool) – If True, show static fields of drinx dataclasses in declaration order, interleaved with dynamic fields.

Return type:

str

Returns:

A multi-line string with the tree diagram.