Feature walkthroughs

• Jit-aware context variables
  • A primer on context variables (sans jit-awareness)
  • Why are context variables useful for JAX programming?
  • Jit-awareness
  • Plumb and transform
  • Recipes
    • Auto-plumbing random keys
    • Stateful nodes in pytrees
  • More specifications
  • Limitations of context variables
• Hello, treepie nodes
  • Dataclasses
  • Frozen sequences
  • Frozen mappings
• Composite nodes
  • What are composite nodes?
  • Examples
    • Example 1
    • Example 2
  • Accessing a child node without unwrapping it
  • Adding tags and properties
  • FrozenSeq and FrozenMap also unwrap composite nodes
• Ancestry aware nodes
  • Quick overview
  • Example
  • A quick use case: weight sharing
  • Ancestry semantics and treepie’s implementation
  • < ONLY WRITTEN UPTO HERE >
  • What does this mean in practice