Outstanding Primitive Rules for Higher Order Automatic Differentiation (jet)

See original GitHub issue

zeros_like
id
add_any
device_put
neg
sign
floor
ceil
round
is_finite
exp
log
expm1
log1p
tanh
sin
cos
atan2
regularized_incomplete_beta
lgamma
igamma
igammac
bessel_i0e
bessel_i1e
erf
erfc
erf_inv
real
imag
complex
conj
abs
sqrt
rsqrt
pow
and
or
xor
add
sub
mul
safe_mul
div
rem
max
min
shift_left
shift_right_arithmetic
shift_right_logical
eq
ne
ge
gt
le
lt
convert_element_type
bitcast_convert_type
conv_general_dilated
dot_general
broadcast
broadcast_in_dim
clamp
concatenate
pad
reshape
rev
transpose
select
slice
dynamic_slice
dynamic_update_slice
gather
scatter-add
scatter
reduce_sum
reduce_prod
reduce_max
reduce_min
reduce_window_sum
reduce_window_max
reduce_window_min
select_and_scatter_add
select_and_gather_add
sort
sort_key_val
tie_in
stop_gradient
while
cond
scan
custom_linear_solve
fft
psum
ppermute
jit(arcsinh,static_argnums=())
cholesky
eigh
triangular_solve
lu
qr
svd
jit(slogdet, static_argnums=())
logit
expit
log_ndtr
random_gamma
asinh
acosh
atanh

@jacobjinkelly @duvenaud @mattjj

Here’s the list of all the ad.jvp_primitives and probably reflects what we’d need for coverage by jet.

Copied from #2363 since that was more about merging into master!

Issue Analytics

State:
Created 4 years ago
Reactions:4
Comments:23 (20 by maintainers)

Top GitHub Comments

1reaction

jacobjinkellycommented, Feb 16, 2022

@pnkraemer Thanks for going through and making this list! I can confirm that the jet rule for expit is implemented and works based on these tests: https://github.com/google/jax/blob/main/tests/jet_test.py#L278-L281

0reactions

mattjjcommented, Dec 21, 2022

@dsheldon thanks for the tips!

I remember trying polynomial composition, but I thought it didn’t give the correct answers (even though I thought it should work), and I never figured out why! I think I messed up _series_to_taylor and _taylor_to_series (referring to the names in your PR). It’s great if this works generically!

Top Results From Across the Web

jax.experimental.jet module - JAX documentation

Jet is an experimental module for higher-order automatic differentiation that does not rely on repeated first-order automatic differentiation.

Higher Order Automatic Differentiation of ...

Abstract. We present semantic correctness proofs of automatic differentiation (AD). We consider a forward-mode AD method on a higher order ...

Automatic differentiation

In mathematics and computer algebra, automatic differentiation (AD), also called algorithmic differentiation, computational differentiation, ...

Taylor-Mode Automatic Differentiation for Higher-Order ...

In defining the higher-order partial derivative rules for each primitive some evaluations can be shared. In computing the total derivative of function.

Automatic Differentiation in PCF

So our result may be seen as an extension Joss's theorem in several directions: to a higher-order language; to a wider set of...