Support all JAX primitives
mstreeter opened this issue · comments
- abs
- acos
- acosh
- 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
- broadcast_in_dim
- cbrt
- ceil
- clamp
- clz
- complex
- concatenate
- cond
- conj
- conv_general_dilated
- convert_element_type
- copy
- cos
- cosh
- create_token
- cumlogsumexp
- cummax
- cummin
- cumprod
- cumsum
- custom_linear_solve
- device_put
- digamma
- div
- dot_general
- dynamic_slice
- dynamic_update_slice
- eq
- erf
- erf_inv
- erfc
- exp
- expm1
- fft
- floor
- gather
- ge
- gt
- igamma
- igamma_grad_a
- igammac
- imag
- infeed
- integer_pow
- iota
- is_finite
- le
- lgamma
- log
- log1p
- logistic
- lt
- max
- min
- mul
- ne
- neg
- nextafter
- not
- or
- outfeed
- pad
- pmax
- pmin
- population_count
- pow
- ppermute
- psum
- random_gamma_grad
- real
- reduce
- reduce_and
- reduce_max
- reduce_min
- reduce_or
- reduce_precision
- reduce_prod
- reduce_sum
- reduce_window
- reduce_window_max
- reduce_window_min
- reduce_window_sum
- reduce_xor
- regularized_incomplete_beta
- rem
- reshape
- rev
- rng_bit_generator
- rng_uniform
- round
- 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
- sqrt
- squeeze
- stop_gradient
- sub
- tan
- tanh
- top_k
- transpose
- while
- xor