OpTree: Optimized PyTree Utilities

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

• PyTree Node Registry
• Integration with functools
• Typing Support
• API References

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

License

OpTree is released under the Apache License 2.0.

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