OpTree: Optimized PyTree Utilities

• Tree Operations
  • Constants
  • Tree Manipulation Functions
  • Tree Reduce Functions
  • PyTreeSpec Functions

• Simplified PyTree Utilities
• PyTreeSpec Constructor API
• PyTree Node Registration
• Integration with dataclasses
• Integration with functools
• Typing Support
• API References

• Integrations with Third-Party Libraries
  • Integration for attrs
  • Integration for JAX
  • Integration for NumPy
  • Integration for PyTorch

License

OpTree is released under the Apache License 2.0.

OpTree is based on JAX’s implementation of the PyTree utility, with significant refactoring and several improvements. The original licenses can be found at JAX’s Apache License 2.0 and Tensorflow’s Apache License 2.0.