JAX Scaled Operators coverage

Summary of JAX LAX operators supported in scalify graph transformation.

JAX LAX operations

Operation	Supported	Remarks
abs	:white_check_mark:
add	:white_check_mark:
acos	:x:
approx_max_k	:x:
approx_min_k	:x:
argmax	:white_check_mark:
argmin	:white_check_mark:
asin	:x:
atan	:x:
atan2	:x:
batch_matmul	:x:
bessel_i0e	:x:
bessel_i1e	:x:
betainc	:x:
bitcast_convert_type	:white_check_mark:
bitwise_not	:x:
bitwise_and	:x:
bitwise_or	:x:
bitwise_xor	:x:
population_count	:x:
broadcast	:white_check_mark:
broadcast_in_dim	:white_check_mark:
cbrt	:x:
ceil	:x:
clamp	:x:
collapse	:white_check_mark:
complex	:x:
concatenate	:white_check_mark:
conj	:x:
conv	:white_check_mark:
convert_element_type	:white_check_mark:
conv_general_dilated	:white_check_mark:
conv_transpose	:white_check_mark:
cos	:white_check_mark:
cosh	:x:
cummax	:x:
cummin	:x:
cumprod	:x:
cumsum	:x:
digamma	:x:
div	:white_check_mark:
dot	:white_check_mark:
dot_general	:white_check_mark:	Limited set of configurations. See below.
dynamic_slice	:x:
dynamic_update_slice	:x:
eq	:white_check_mark:
erf	:x:
erfc	:x:
erf_inv	:x:
exp	:white_check_mark:
expand_dims	:white_check_mark:
expm1	:x:
fft	:x:
floor	:x:
full	:question:
full_like	:question:
gather	:x:
ge	:white_check_mark:
gt	:white_check_mark:
igamma	:x:
igammac	:x:
imag	:x:
index_in_dim	:x:
index_take	:x:
iota	:white_check_mark:
is_finite	:white_check_mark:
le	:white_check_mark:
lt	:white_check_mark:
lgamma	:x:
log	:white_check_mark:
log1p	:x:
logistic	:x:
max	:white_check_mark:
min	:white_check_mark:
mul	:white_check_mark:
ne	:white_check_mark:
neg	:white_check_mark:
nextafter	:x:
pad	:white_check_mark:
polygamma	:x:
pow	:x:
real	:x:
reciprocal	:x:
reduce	:white_check_mark:
reduce_precision	:white_check_mark:
reduce_window	:white_check_mark:
reshape	:white_check_mark:
rem	:x:
rev	:white_check_mark:
round	:x:
rsqrt	:x:
scatter	:x:
scatter_add	:x:
scatter_max	:x:
scatter_min	:x:
scatter_mul	:x:
select	:white_check_mark:
shift_left	:x:
shift_right_arithmetic	:x:
shift_right_logical	:x:
slice	:white_check_mark:
slice_in_dim	:white_check_mark:
sign	:x:
sin	:white_check_mark:
sinh	:x:
sort	:x:
sort_key_val	:x:
sqrt	:x:
square	:x:
squeeze	:white_check_mark:
sub	:white_check_mark:
tan	:x:
tie_in	:x:	Deprecated in JAX
top_k	:x:
transpose	:white_check_mark:
zeta	:x: