Overview of implemented operations

[sebffischer]

The stableHLO specification includes the operations listed below, many of which are already implemented.

In general, when implementing an operation, we:

1. don't support complex types.
2. don't handle quantized types.

The main goal when implementing a new operation is to implement the type inference according to the rules from SPEC.md.
Ideally, every check should be annotated with the corresponding criterion ((C1), (C2), ...).

• all_gather (sharding, multi process grid)
• all_reduce (sharding, multi process grid)
• all_to_all (sharding, multi process grid)
• uniform_dequantize (quantization not supported yet)
• uniform_quantize (quantization not supported yet)
• collective_broadcast (distributed)
• collective_permute (distributed)
• infeed (distributed)
• outfeed (distributed)
• recv (distributed)
• send (distributed)
• replica_id (multi process)
• partition_id (multi process)
• tuple (is being deprecated soon)
• map (is being deprecated)
• get_tuple_element (tuple is being deprecated soon)
• rng (non-reproducible rngs are not super useful)
• complex (complex not supported yet)
• real (complex not supported yet)
• imag (complex not supported yet)
• batch_norm_grad (not sure if needed)
• batch_norm_inference (not sure if needed)
• batch_norm_training (not sure if needed)
• dynamic_broadcast_in_dim
• dynamic_conv
• dynamic_gather
• dynamic_iota
• dynamic_pad
• dynamic_reshape
• dynamic_slice
• dynamic_update_slice
• get_dimension_size
• composite
• fft
• optimization_barrier
• select_and_scatter
• reduce_scatter
• convolution
• reduce_precision
• gather
• scatter
• reduce_window
• triangular_solve
• count_leading_zeros
• iota
• custom_call
• pad
• return
• abs
• add
• after_all
• and
• atan2
• bitcast_convert
• broadcast_in_dim
• case
• cbrt
• ceil
• cholesky
• clamp
• compare
• concatenate
• constant
• convert
• cosine
• divide
• dot_general
• exponential
• exponential_minus_one
• floor
• if
• is_finite
• log
• log_plus_one
• logistic
• maximum
• minimum
• multiply
• negate
• not
• or
• popcnt
• power
• reduce
• remainder
• reshape
• reverse
• rng_bit_generator
• round_nearest_afz
• round_nearest_even
• rsqrt
• select
• shift_left
• shift_right_arithmetic
• shift_right_logical
• sign
• sine
• slice
• sort
• sqrt
• subtract
• tan
• tanh
• transpose
• while
• xor

[sebffischer]

Next simple ops to look at:

• sort
• not, shift_left, shift_right_arithmetic, shift_right_logical: use R functions bitwNot, ... for testing. If tricky, ignore
• is_finite
• popcnt
• reverse: look at broadcast_dims of BroadcastInDim, should hopefully 'just work'
• transpose, similar to reverse (check BroadcastInDim) (I am doing this)
• reshape (output shape is custom_attr, also look at BroadcastInDim 'output_shape' custom_attr.
• round_nearest_afz
• select
• slice
• cholesky
• concatenate

Afterwards we can look at other important ops:

• while
• map (might be deprecated)
• case (similar to if)
• reduce (already done)