Currently hlo-legalize-to-linalg on works with mhlo.scatter when the body only has a single op https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc#L2809. Correct me if I’m wrong, but this seems to only correspond to Identity for the update_computation argument in XLA since a return op is always required. Is anyone working on extending this a bit more? Are there any other ways to get this lowered?