ASEM000/PyTreeClass: v0.2.0

Mostly rewritten from scratch.

• Drop the standard library dataclasses.dataclass for faulty logic treating jax.numpy as mutable object from python 3.11 see: #33

• Implements a simpler dataclass-like version that additionally implements:

  • pos_only to mark the argument as positional only
  • callbacks to apply a set of functions (e.g., validators/converters) on input see: #37

• Drop metadata-based approach for marking fields static. move to a transparent wrapper approach for faster, simpler, and flexible marking of leaves static see here

• tree_summary, tree_diagram, tree_str, and tree_repr can now work on any pytree(including registered pytrees) out of the box see
• Remove tree.summary() in favor of pytc.tree_summary(model)
• Remove tree.tree_diagram() in favor of pytc.tree_diagram(model)
• Remove tree.tree_box()
• Add tree_{repr,str,diagram,summary,mermaid}(..., depth) to control cutoff depth condition for visualization.
• Make math operations on leaves opt-in feature via treeckass(... , leafwise=True)
• Drop filtering by metadata , types, and name. Instead, use tree_map_with_trace to access leaves' name,type,index, and metadata at each level along their path. see #35
• Add bcmap : a function transformation for automatic broadcasting. see readme More section.
  • using bcmap with treeclass(... , leafwise=True) achieves a data model that enables applying numpy functions like numpy.where on arbitrary PyTrees without the need for tree_map.
• Faster flatten/unflatten see benchmarks comparison with flax/equinox.

Internally, unifying indexing and viz tools by implementing custom jax-like registry to record names, types, indices, and metadata rules for common python data structures and treeclass wrapped class. You can check the readme More section for registering custom classes to work fully with the viz/indexing tools similar to Jax registration process.