[Feature request] prospects of integrating jax's jacfwd, jacrev, and hessian functions in mlx

Open LiPingYen opened this issue 2 years ago • 9 comments

Discussed in https://github.com/ml-explore/mlx/discussions/154

^{Originally posted by LiPingYen December 14, 2023} In JAX, there are jacfwd, jacrev, and hessian functions for transforming the objective function into functions that compute first-order or second-order derivatives. I'm curious to know if MLX has plans to incorporate these three functions in the future.

Dec 14 '23 10:12 LiPingYen

Definitely would like to have these transforms. I will make it a feature request. Contributions for this are very welcome until we have the bandwidth to get to it.

Dec 14 '23 15:12 awni

Hi @awni ,

I appreciate your responsiveness. Look forward to leveraging the capabilities of XLM to execute jacfwd, jacrev, and hessian functionalities in the near future.

Dec 14 '23 15:12 LiPingYen