transpose_rules for nonlinear primitives.

[mavenlin]

Here's my understanding of transpose_rules.

Transpose rules are registered for linear primitives. When a function is linearized with jvp, the linearized function can be transposed op by op, and that's how vjp is computed in jax.

1. Is vjp the only transform that calls the transpose rules under the hood? Any example that it is used in other types of transforms?

2. Say if I try to implement a custom primitive invoke_p, with invoke_impl = lambda x, f: f(x). Here f is a parameter of the primitive, which itself is a python function. Therefore, whether the output of invoke is linear to x depends on f.

Now since invoke is a primitive, I need to define jvp and transpose rules for it. jvp rule is relatively easy, because it requires linearization wrt x, so as long as f is linearizable we're fine.

When it comes to transpose rules, the question comes, should we even implement transpose rule for this primitive? As it is not guaranteed to be linear. Or should we still give it a transpose rule by linearizing f within the transpose rule?

Thanks for reading my question. The above may sound wierd requirement for a primitive, but it could happen when we try to make the differentiation engine work for some custom calls.

[froystig]

Is vjp the only transform that calls the transpose rules under the hood? Any example that it is used in other types of transforms?

In addition to VJP, there is also the more direct jax.linear_transpose, which assumes linearity of the input function.

When it comes to transpose rules, the question comes, should we even implement transpose rule for this primitive? As it is not guaranteed to be linear.

Linearity only matters for primitives bound under the tangent computation of a JVP (i.e. during linearization), or for primitives bound under jax.linear_transpose. If you never expect to transpose invoke_p because it is never bound in a tangent computation, and because you don't plan to support linear_transpose, then you don't need a transposition rule.

Or should we still give it a transpose rule by linearizing f within the transpose rule?

You could give invoke_p a transpose rule that transposes f using jax.linear_transpose. When linear_tranpose executes, it will error if f is not structurally linear. Presumably, if you've bound invoke_p in the tangent computation of a JVP, it is linear and you wouldn't expect this error!