return true;
      };

  return RewriteOp(op, rewriter, get_broadcasted_shape);
}

template <typename Op>
LogicalResult ConvertResultsBroadcastableShapeOp::RewriteEqOp(
    Operation* op, PatternRewriter& rewriter) const {
  auto eq_op = llvm::dyn_cast_or_null<Op>(op);
  if (eq_op && eq_op.getIncompatibleShapeError())
    return RewriteOp(op, rewriter, OpTrait::util::getBroadcastedShape);

#include "tensorflow/core/framework/op_kernel.h"

namespace mlir {
namespace mhlo {

class Tf2XlaRewriterTestPeer;

class Tf2XlaRewriter {
 public:
  static mlir::LogicalResult RewriteOp(mlir::Operation* op,
                                       mlir::PatternRewriter& rewriter,
                                       const std::string& device_type);

 private:
  friend class Tf2XlaRewriterTestPeer;

          (prefer_tf2xla_ &&
           IsOpAllowedTf2xlaPreferred(abstractOp->getTypeID())))) {
      return failure();
    }

    return Tf2XlaRewriter::RewriteOp(op, rewriter, device_type_);
  }

 private:
  std::string device_type_;
  bool prefer_tf2xla_;
  bool use_tf2xla_hlo_importer_;
};

bool ShouldRefineTypeTo(Type original_ty, Type updated_ty) {

    OpBuilder op_builder(&op);
    EmptyPatternRewriter pattern_rewriter(op_builder);

    LogicalResult result =
        Tf2XlaRewriter::RewriteOp(&op, pattern_rewriter,
                                  /*device_type=*/"XLA_CPU_JIT");
    if (!result.succeeded()) {
      return tsl::errors::Internal("Failed to rewrite op");
    }

      hlo_module.entry_computation()->root_instruction();

  return root_instruction->opcode() == xla::HloOpcode::kTuple;
}

};  // namespace

LogicalResult Tf2XlaRewriter::RewriteOp(Operation* op,
                                        PatternRewriter& rewriter,
                                        const std::string& device_type) {
  Tf2XlaRewriter tf2xla_rewriter(op, rewriter, device_type);