}

//===----------------------------------------------------------------------===//
// AllToAll op legalizations.
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: func @alltoall_basic
// See https://www.tensorflow.org/api_docs/python/tf/raw_ops/AllToAll
func.func @alltoall_basic(%input: tensor<1x2xf32>) -> tensor<2x1xf32> {

*   Fixes a null pointer exception when `Exit` node is not preceded by `Enter`
    op
    ([CVE-2021-41217](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-41217))
*   Fixes an integer division by 0 in `tf.raw_ops.AllToAll`
    ([CVE-2021-41218](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-41218))
*   Fixes a use after free and a memory leak in `CollectiveReduceV2`

    Res<TF_BoolTensor, [{The reduced tensor.}]>:$output
  );

  TF_DerivedOperandTypeAttr Tidx = TF_DerivedOperandTypeAttr<1>;

  let hasVerifier = 1;
}

def TF_AllToAllOp : TF_Op<"AllToAll", [Pure, TF_NoConstantFold]> {
  let summary = "An Op to exchange data across TPU replicas.";

  let description = [{
On each replica, the input is split into `split_count` blocks along