XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(                                                     \
      Name("XlaBroadcastHelper").HostMemory("broadcast_dims").Device(DEVICE),  \
      XlaCompileOnDemandOp);                                                   \
  REGISTER_KERNEL_BUILDER(Name("XlaSelfAdjointEig").Device(DEVICE),            \

  %0 = "tf.Const"() {value = dense<2> : tensor<1xi64>} : () -> tensor<1xi64>
  // expected-error @+2 {{'tf.XlaBroadcastHelper' op failed to infer returned types}}
  // expected-error @+1 {{broadcast_dims must have size equal to the smaller argument rank}}
  %lhs_output, %rhs_output = "tf.XlaBroadcastHelper"(%arg0, %arg1, %0) : (tensor<2x3x5xi32>, tensor<5x2xi32>, tensor<1xi64>) -> (tensor<2x3x5xi32>, tensor<2x1x5xi32>)
  func.return
}

      output_memory_types[i] = tensorflow::HOST_MEMORY;
    }
  }
  return output_memory_types;
}

static auto const ops_triggering_xla_compilation =
    new absl::flat_hash_set<std::string>{"XlaBroadcastHelper",
                                         "XlaCallModule",
                                         "XlaConv",
                                         "XlaConvV2",

      "UnsortedSegmentProd",
      "UnsortedSegmentSum",
      "VarIsInitializedOp",
      "VariableShape",
      "Where",
      "While",
      "XlaAllReduce",
      "XlaBroadcastHelper",
      "XlaCallModule",
      "XlaConcatND",
      "XlaConv",
      "XlaConvV2",
      "XlaCustomCall",
      "XlaCustomCallV2",
      "XlaDequantize",
      "XlaDot",

    TensorOf<[TF_Bfloat16, TF_Float16, TF_Float32, TF_Int32, TF_Uint32]>:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

def TF_XlaBroadcastHelperOp : TF_Op<"XlaBroadcastHelper", [InferTensorType, Pure, TF_NoConstantFold]> {
  let summary = "Helper operator for performing XLA-style broadcasts";

  let description = [{