ML Katas

by: newest upvotes saves

Custom Gradient with `jax.custom_vjp`

yesterday hard (<30 mins) | autograd jax custom-gradient vjp

Implement a function with a custom gradient using `jax.custom_vjp`. This is useful for numerical stability or for defining gradients for non-differentiable operations. A good example is a function...

1