Support all JAX primitives

[mstreeter]

• [X] abs
• [ ] acos
• [ ] acosh
• [X] add
• [ ] after_all
• [ ] all_gather
• [ ] all_to_all
• [ ] and
• [ ] approx_top_k
• [ ] argmax
• [ ] argmin
• [ ] asin
• [ ] asinh
• [ ] atan
• [ ] atan2
• [ ] atanh
• [ ] axis_index
• [ ] bessel_i0e
• [ ] bessel_i1e
• [ ] bitcast_convert_type
• [X] broadcast_in_dim
• [x] cbrt
• [ ] ceil
• [ ] clamp
• [ ] clz
• [ ] complex
• [ ] concatenate
• [ ] cond
• [ ] conj
• [X] conv_general_dilated
• [X] convert_element_type
• [ ] copy
• [ ] cos
• [ ] cosh
• [ ] create_token
• [ ] cumlogsumexp
• [ ] cummax
• [ ] cummin
• [ ] cumprod
• [ ] cumsum
• [ ] custom_linear_solve
• [ ] device_put
• [ ] digamma
• [X] div
• [X] dot_general
• [ ] dynamic_slice
• [ ] dynamic_update_slice
• [ ] eq
• [ ] erf
• [ ] erf_inv
• [ ] erfc
• [X] exp
• [ ] expm1
• [ ] fft
• [ ] floor
• [ ] gather
• [ ] ge
• [ ] gt
• [ ] igamma
• [ ] igamma_grad_a
• [ ] igammac
• [ ] imag
• [ ] infeed
• [X] integer_pow
• [X] iota
• [ ] is_finite
• [ ] le
• [ ] lgamma
• [X] log
• [ ] log1p
• [x] logistic
• [ ] lt
• [ ] max
• [ ] min
• [X] mul
• [ ] ne
• [X] neg
• [ ] nextafter
• [ ] not
• [ ] or
• [ ] outfeed
• [ ] pad
• [ ] pmax
• [ ] pmin
• [ ] population_count
• [X] pow
• [ ] ppermute
• [ ] psum
• [ ] random_gamma_grad
• [ ] real
• [ ] reduce
• [ ] reduce_and
• [ ] reduce_max
• [ ] reduce_min
• [ ] reduce_or
• [ ] reduce_precision
• [ ] reduce_prod
• [X] reduce_sum
• [ ] reduce_window
• [ ] reduce_window_max
• [ ] reduce_window_min
• [ ] reduce_window_sum
• [ ] reduce_xor
• [ ] regularized_incomplete_beta
• [ ] rem
• [X] reshape
• [ ] rev
• [ ] rng_bit_generator
• [ ] rng_uniform
• [ ] round
• [x] rsqrt
• [ ] scan
• [ ] scatter
• [ ] scatter-add
• [ ] scatter-max
• [ ] scatter-min
• [ ] scatter-mul
• [ ] select_and_gather_add
• [ ] select_and_scatter
• [ ] select_and_scatter_add
• [ ] select_n
• [ ] shift_left
• [ ] shift_right_arithmetic
• [ ] shift_right_logical
• [ ] sign
• [ ] sin
• [ ] sinh
• [ ] slice
• [ ] sort
• [x] sqrt
• [X] squeeze
• [ ] stop_gradient
• [X] sub
• [ ] tan
• [ ] tanh
• [ ] top_k
• [X] transpose
• [ ] while
• [ ] xor