Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

Please consider lifting custom_vjp

See original GitHub issue

According to the documentation, there is no flax.custom_vjp yet. Are there any plans to implement it?

I love the design of flax.vmap, which adds variable_axes and split_rngs. If a similar approach were taken with custom_vjp, then the ability to lift parameters that I requested https://github.com/google/flax/issues/1712 might be unnecessary.

Issue Analytics

State:
Created 2 years ago
Comments:5 (3 by maintainers)

Top GitHub Comments

1reaction

levskayacommented, Dec 31, 2021

yeah no worries, let us know if it works for you.

0reactions

jheekcommented, Jan 3, 2022

Lol I linked the PR to itself instead of to this issue. Thanks for mentioning @levskaya

Top Results From Across the Web

Tracer seems to leak into vjp through nondiff_argnums #2912

I think this is a bug in how I implemented nondiff_argnums : I made it equivalent to lexical closure, but since custom_vjp functions...

Custom derivative rules for JAX-transformable Python functions

A related application is to enforce a differentiation convention, perhaps at a boundary. Consider the function ...

The L.I.F.T. Lounge®️ Delray Beach FL Private Training Studio

VIP Private Training Sessions at The L.I.F.T. Lounge®, located in our brand new studio in Delray Beach, Florida using our patent-pending custom equipment ......

Lift Tickets - Mountain High Resorts

Ski or Snowboard for just $79 per Day, $39 per Night (5pm-10pm). Please purchase your discounted tickets at the VIP Window at the...

Ep. #929: Greg Everett on How to Start Olympic Weightlifting ...

I think you're going to enjoy it! Timestamps: 0:00 – Legion VIP One-on-One Coaching: https://www.muscleforlife.show ...