operations.md

TessellateIPU supported operations

JAX LAX operations

Note: Inplace operations are represented by additional primitives (e.g. add_inplace_p). For binary operations, only the first argument is supported for inplace update.

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

Limitations

• dot_general: Only supporting 16x16 right hand side at the moment.
• gather/scatter: Only first axis gather/scatter supported.

Additional IPU optimized operations

Operation	Supported	Inplace	Remarks
scaled_add	✅	❌	Fused scaled add
scaled_sub	✅	❌	Fused scaled sub