//   transpose(updates, permutation={3,4,0,1,2})
//
// Note: NormalizeIndexVector is assumed to have run on the indices already so
// that the index_vector_dim is the trailing dimension in `indices`.
LogicalResult CanonicalizeScatterUpdates(
    Operation* scatter_op, llvm::ArrayRef<int64_t> update_window_dims,
    const Value& indices, const ShapedType& indices_type, Value& updates,
    ShapedType& updates_type, ConversionPatternRewriter& rewriter);

#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/util.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {

LogicalResult CanonicalizeScatterUpdates(
    Operation* scatter_op, llvm::ArrayRef<int64_t> update_window_dims,
    const Value& indices, const ShapedType& indices_type, Value& updates,
    ShapedType& updates_type, ConversionPatternRewriter& rewriter) {