jet of pow using comp with exp, mul, log

[jacobjinkelly]

This implementation is numerically unstable (the test begins to fail when the upper limit is about 35), but it's unclear if there's another way to implement jet of pow easily.

Co-authored-by: Jesse Bettencourt jessebett@cs.toronto.edu (@jessebett)
Co-authored-by: David Duvenaud duvenaud@cs.toronto.edu (@duvenaud)

[googlebot]

All (the pull request submitter and all commit authors) CLAs are signed, but one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project. Please have them confirm that by leaving a comment that contains only @googlebot I consent. in this pull request.

Note to project maintainer: There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent. In those cases, you can manually confirm consent of the commit author(s), and set the cla label to yes (if enabled on your project).

ℹ️ Googlers: Go here for more info.

[shoyer]

If you're going to define things this way, I'd just call back into jet:

Suggested change

	u_, r_ = primals_in
	u_terms, r_terms = series_in
	u = [u_] + u_terms
	r = [r_] + r_terms
	logu_, logu_terms = _log_taylor((u_, ), (u_terms, ))
	v_, v_terms = _bilinear_taylor_rule(lax.mul_p, (logu_, r_), (logu_terms, r_terms))
	v_, v_terms = _exp_taylor((v_, ), (v_terms, ))
	return v_, v_terms
	return jet(lambda x, y: lax.exp(y * lax.log(x), primals_in, series_in)

[duvenaud]

Haha, I feel silly now. We started by trying to derive a giant monolithic formula for pow(x, y), but we already had all the ingredients!

[jacobjinkelly]

@shoyer Thanks for pointing this out! I realized that I should change the primal_out propagated through the exp rule, so I incorporated your suggestion for propagation through mul and log in 77b7d9b (added you as a co-author)

[jessebett]

@shoyer yeah! We're still unclear if there is a benefit from writing the expanded rules for primitives that can be composed of other primitives vs just calling jet on that composition like you have here.

[shoyer]

I guess the more numerically stable but also more expensive alternate would be to recursively apply the JVP rule?

[jessebett]

@shoyer yeah you can see in the tests the function that looks like recursive calls to jvp:
https://github.com/google/jax/blob/c9c14c02d4c0e2ed6ebbc9bb4e6071a9d1c548fc/tests/jet_test.py#L33-L43

However this is significantly slower. Interesting to see if it's more numerically stable, though, because I'd expect the opposite. Not confident about intuition.

[jacobjinkelly]

I guess the more numerically stable but also more expensive alternate would be to recursively apply the JVP rule?

I think 77b7d9b fixes this. I changed the primal_out propagated through the exp rule to be x ** y (before it was exp(y * log x).

I modified the tests in c75855b, to include the option to check jets on only finite values (by casting np.inf, -np.inf, and np.nan to np.nan). I also added separate lims (to keep the base x positive but allow the power y to be negative). With these modifications the tests for jet(pow) pass on a wide range of values.

[shoyer]

Looks good to me!

@jacobjinkelly That's kind of you to mark me as a co-author, but I think a one line review suggestion barely qualifies :)

@jessebett @duvenaud googlebot still needs your OK on this (@googlebot I consent.) before we can merge

[duvenaud]

@googlebot I consent.

[googlebot]

We found a Contributor License Agreement for you (the sender of this pull request), but were unable to find agreements for all the commit author(s) or Co-authors. If you authored these, maybe you used a different email address in the git commits than was used to sign the CLA (login here to double check)? If these were authored by someone else, then they will need to sign a CLA as well, and confirm that they're okay with these being contributed to Google.
In order to pass this check, please resolve this problem and then comment @googlebot I fixed it.. If the bot doesn't comment, it means it doesn't think anything has changed.

ℹ️ Googlers: Go here for more info.

[shoyer]

@jacobjinkelly could you kindly either remove my "co-author" status from the commit or use by google email instead (shoyer@google.com)?

[googlebot]

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

[jacobjinkelly]

@shoyer Ok rebased the relevant commit with your google email!

[shoyer]

@jessebett this still needs your consent

[jessebett]

@googlebot I consent.

[shoyer]

Thanks everyone!

[jacobjinkelly]

just to keep track this addresses #2431