Redundant calculation of derivatives of power function

[titaneric]

Hi,

It seems that this repo have the same issue as HIPS/autograd#541 and pytorch/pytorch#28651 .

However, it's quiet more complicated than I fix in autograd, so I have no clues how to fix it.

The following is here I found

• The min and max function will take the ans into account to determine their derivatives, and the type of these two is standard_binop, which is same as power.
• The definition of VJP in jax is different than autograd, which is more dynamic ( I mean the argument of VJP can be different, not the fixed form lambda g, ans, x, y )

Could your priors guide me to fix this issue?

It may be solved after I finally find out how jax works, but if someone can help me or give me some hints, it will be great!

[hawkinsp]

If you want access to ans from a JVP rule, you can use ad.defjvp2 instead of ad.defjvp. The former provides ans as the second argument to the JVP rule. We should document these a bit better.

In this case, you could certainly use ans for the rule _pow_jvp_rhs.

(Although note that under a jit XLA will perform common subexpression elimination, so this kind of optimization doesn't matter quite as much as it does in Autograd. But it'd still be a reasonable thing to do.)

[titaneric]

I did not notice that the difference between ad.defjvp and ad.defjvp2, thanks for pointing it out.

I fix it, hope that it help.