Add ExtractDiag, Cholesky, and Solve Ops to jaxification

[twiecki]

Closes #49.

[brandonwillard]

Looks good, just needs some tests.

[twiecki]

Thanks, how do I handle kwargs?

[brandonwillard]

Thanks, how do I handle kwargs?

Which kwargs? For the jax functions?

In general, the jax function arguments are going to come from either the arguments to the JAX closure/function constructed in the Op dispatch function (e.g. jax_funcify_ExtractDiag)—which themselves are determined by the arguments passed to the inputs argument of the Op's perform method—or from fields on the Op instance itself. In the latter case, those fields simply need to be added to the constructed closure/function via reference to the original Op object or—more preferably—by use of default keyword values in the closure/function (e.g. def jax_funcify_ExtractDiag(..., axis=op_axis_value)) where op_axis_value = op.axis preceding the function definition).

[brandonwillard]

Just pushed some changes. Needs a few more tests, though.

[twiecki]

Added more tests, I think this covers everything. Not sure why the lower=False doesn't work.

[twiecki]

OK, figured out lower.

[twiecki]

Dang, messed up some commits when rebasing. Closing in favor of #59.

[brandonwillard]

I think it would be easier if we just wait for the test on this one to finish, merge it, then rebase #59; especially since #59 is behind this branch.