https://github.com/mfschubert/mewtax
Differentiable minimization in jax using Newton's method
https://github.com/mfschubert/mewtax
jax optimization
Last synced: over 1 year ago
JSON representation
Differentiable minimization in jax using Newton's method
- Host: GitHub
- URL: https://github.com/mfschubert/mewtax
- Owner: mfschubert
- License: mit
- Created: 2024-01-11T17:08:29.000Z (over 2 years ago)
- Default Branch: main
- Last Pushed: 2024-06-14T19:53:05.000Z (about 2 years ago)
- Last Synced: 2025-02-15T12:47:45.983Z (over 1 year ago)
- Topics: jax, optimization
- Language: Python
- Homepage:
- Size: 27.3 KB
- Stars: 1
- Watchers: 1
- Forks: 0
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Differentiable minimization in jax using Newton's method
`v0.0.1`
This project essentially repackages code from the [implicit layers tutorial](https://implicit-layers-tutorial.org/implicit_functions/) to provide a `minimize_newton` function.
Given a function `fn(params, z)`, it finds the `z_star` which minimizes `fn` for given `params`. Further, the gradient of the solution with respect to `params` can be computed; this is done using a custom vjp rule, as shown in the tutorial.
## Installation
mewtax can be installed via pip:
```
pip install mewtax
```