Best way to define gradient for linear black box function

Dear jax team,

I’d like to use a black box function in jax with grad where the function is a linear operator, i.e. a function f(x) executing A.dot(x) where A is a matrix too large for memory so A is computed on-the-fly.

I guess I could define its gradient with jax.custom_gradient, but I assumed there are performance benefits in telling jax that the operation is linear? Either way, I would have to provide another black box function computing A.T.dot(x).

Is there a way in jax to do this? Or should I just provide the gradient rule via jax.custom_gradient?

[shoyer]

Are the non-zero entries of the sparse matrices all in memory at some point? Or can the operators be totally opaque black-box functions?

The idea is make everything work based on black-box functions for computing matrix-vector products.

(We don’t have support for explicit sparse matrices in JAX yet.)

[jekbradbury]

If you’re willing to define your function as a core.Primitive, you can set it up for AD using ad.deflinear(primitive, transpose_rule), where transpose_rule is a function that behaves like A.T.dot(x). This notebook may be helpful for elucidating the current (internal/unstable) API surface for primitives.