// Training with non-broadcastable shape
  %cst = arith.constant dense<0.0> : tensor<4xf32>

  %5 = "tf.AddV2"(%4, %1): (tensor<4x4xf32>, tensor<1xf32>) -> tensor<4x4xf32>
  %6 = "tf.Log"(%5): (tensor<4x4xf32>) -> tensor<4x4xf32>

  // This is a legal canonicalization because constant shape 4xf32 is
  // broadcastable to 4x4xf32, however we currently do not support this case,
  // and canonicalize only if the constant is a scalar.
  // CHECK: %[[ADD2:.*]] = "tf.AddV2"
  // CHECK: %[[LOG2:.*]] = "tf.Log"(%[[ADD2]])

    rewriter.replaceOpWithNewOp<SourceOp>(op, result_type, lhs, rhs);
    return success();
  }
};

// This specialization is for TF SelectV2 op. SelectV2 op have three inputs and
// they should have broadcastable shapes.
template <>
class ApplyExplicitBroadcasting<TF::SelectV2Op>
    : public OpRewritePattern<TF::SelectV2Op> {
 public:
  using OpRewritePattern<TF::SelectV2Op>::OpRewritePattern;

  // Check compatibility of batch dimensions if both input shapes are known.
  // BatchMatMul should have exactly the same batch dimensions and
  // BatchMatMulV2 should have broadcastable batch dimensions.
  //
  // The last two dimensions are non-batch dimensions that don't need to
  // participate in batch dimension compatibility check.
  if (std::is_same<OpT, BatchMatMulOp>()) {

  string tflRuntimeDescription = desc;
}

class TFL_OperandsHaveSameShapesOrBroadcastableShape<
    list<int> indices, int max_bcast_rank> :
  TFL_RuntimePredOpTrait<"operands do not have the same shape or "
      "broadcastable shapes within the rank " # max_bcast_rank,
    CPred<"TFL::VerifyOperandsHaveSameShapesOrBroadcastableShape("
            "$_op, llvm::ArrayRef<unsigned>({" # !interleave(indices, ", ") #

                                          Value output, BoolAttr adj_x,
                                          BoolAttr adj_y) {
  // TensorFlow BatchMatMulOp allows the batch dimensions to be broadcastable
  // while the XlaDotV2Op doesn't. So we have to broadcast them beforehand.
  BroadcastBatchDimensionsForBatchMatMul(builder, loc, input, weight);

  // Both input and weight have the same rank after broadcasting.

  return Maximum(scope, ids, ZerosLike(scope, ids));
}

// Helper function for unsorted segment ops.
// Returns a mask of where 'ids' are positive, reshaped so that it will be
// broadcastable to the result shape of gathering params by ids.
Output GetIsPositive(const Scope& scope, const Output& params,
                     const Output& ids) {
  Output is_positive = GreaterEqual(scope, ids, ZerosLike(scope, ids));

LogicalResult NotEqualOp::verify() {
  NotEqualOp op = *this;
  // If we allow inputs to have incompatible type, then nothing to do.
  if (!op.getIncompatibleShapeError()) return success();

  // Otherwise, check inputs are broadcastable.
  return mlir::OpTrait::impl::verifyCompatibleOperandBroadcast(
      op.getOperation());
}

void NotEqualOp::build(OpBuilder &builder, OperationState &result, Value x,

// CHECK-LABEL: greater_equal
// CHECK:  tfl.greater_equal(%arg0, %arg1) : (tensor<8x16xf32>, tensor<8x16xf32>) -> tensor<8x16xi1>
// CHECK:  return
}

//TODO(b/136498739): Add failure test for non-broadcastable types, since currently
// we can't catch this error.
func.func @less_equal(%arg0: tensor<8x16xf32>, %arg1: tensor<8x16xf32>) -> tensor<8x16xi1> {

              'has_bias': [True, False],
              'batch_sizes': [([], []), ([10], [10]), ([2, 3], [2, 3])],
              'target_opset': [quant_opts_pb2.XLA],
          },
          # Test broadcastable batch sizes.
          {
              'activation_fn': [None],
              'has_bias': [True],
              'batch_sizes': [
                  ([2], []),
                  ([], [2]),