Fox is an OCaml library for automatic differentiation and numerical computing, inspired heavily by Jax. It provides a flexible framework for automatically differentiating and XLA-jit-compiling tensor-based computations.

Please note that this library is still under active development and lacks many features.

• Automatic Differentiation: Support for forward-mode and reverse-mode automatic differentiation
• Higher-Order Derivatives: Compute nth-order derivatives of functions
• JIT Compilation: Just-in-time compilation to XLA
• Tree-based Value Representation: Efficient handling of complex data structures through tree-based representations

DYLD_LIBRARY_PATH=/absolute-path-to-ocaml-xla-extension-lib dune build @default @runtest -w

• Support non-singleton tensors in vjp
• Better shape and type story
• Basic tensor operations for simple neural network example
  • 2d matmuls
  • sum
  • random other things
• Print XLA HLO module
• Add JIT caching story
• Pytorch backend
• Testing framework for diffing op backends (XLA, Pytorch, OCaml)
  • Quickcheck generators
  • Diff XLA backend and OCaml
  • Diff OCaml/XLA vjp and Pytorch
• Support various types (ints, bools, etc.)
• Shape inference?
• Mixed-precision support
• Custom operator support