diff --git a/docs/operations.md b/docs/operations.md index d378181..5b35e88 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -2,6 +2,121 @@ ## JAX LAX operations - +| Operation | Supported | Inline | Remarks | +| ------------ | --------- | ------ |-------- | +| `abs` | | | | +| `add` | | | | +| `acos` | | | | +| `approx_max_k` | | | | +| `approx_min_k` | | | | +| `argmax` | | | | +| `argmin` | | | | +| `asin` | | | | +| `atan` | | | | +| `atan2` | | | | +| `batch_matmul` | | | | +| `bessel_i0e` | | | | +| `bessel_i1e` | | | | +| `betainc` | | | | +| `bitcase_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` | | | | +| `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` | | | | +| `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` | | | | +| `top_k` | | | | +| `transpose` | | | | +| `zeta` | | | | ## Additional IPU optimized operations