operators.md

JAX Scaled Operators coverage

Summary of JAX LAX operators supported in scalify graph transformation.

JAX LAX operations

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