Support JAX-like function transforms (grad, vmap), similar to functorch

apache / mxnet

Lightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more

https://mxnet.apache.org

Apache License 2.0

20.78k stars 6.79k forks source link

Support JAX-like function transforms (grad, vmap), similar to functorch #20561

Open learning-chip opened 3 years ago

learning-chip commented 3 years ago

functorch allows function transforms like grad(func)(*inputs), borrowing from JAX. This opens new doors like:

High-order gradient by grad(grad(func))
Second-order quasi-Newton optimizers that need Jacobian as input. Notably (L-)BFGS and its randomized variants.
Better probablistic programming support as in NumPyro

Happy to discuss more if there are any interests.

github-actions[bot] commented 3 years ago

Welcome to Apache MXNet (incubating)! We are on a mission to democratize AI, and we are glad that you are contributing to it by opening this issue. Please make sure to include all the relevant context, and one of the @apache/mxnet-committers will be here shortly. If you are interested in contributing to our project, let us know! Also, be sure to check out our guide on contributing to MXNet and our development guides wiki.