ArrayRef<StateT *> results) {
    if (auto cast = dyn_cast<TF::CastOp>(op)) {
      this->join(results[0], *operands[0]);
    } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {
      for (auto &region : while_op->getRegions()) {
        for (auto [arg, value] :
             llvm::zip(region.getArguments(), while_op->getOperands())) {

    auto logits_softmax = Softmax(scope, logits);

    auto grad_grad_expand = ExpandDims(scope, grad_grad, 1);
    auto logits_softmax_expand = ExpandDims(scope, logits_softmax, 2);
    auto matmul_result =
        BatchMatMul(scope, grad_grad_expand, logits_softmax_expand);
    axis.push_back(1);
    auto squeeze_result = Squeeze(scope, matmul_result, Squeeze::Axis(axis));

  return false;
}

}  // namespace

bool TargetHardware::Init() {
  auto* hardware_ops_factory = GetRegisteredTargetHardwareOps();
  for (auto& hardware_ops : *hardware_ops_factory) {
    if (hardware_ops->hardware_typeid != this->GetTypeId()) continue;
    auto& op_factories = hardware_ops->target_hardware_ops_factory;
    for (auto& op_factory : op_factories) {
      hardware_ops_.emplace_back(op_factory.getSecond()());
    }

  return %2 : tensor<10x10xi8>
})mlir";
  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({10, 10}));
  TF_ASSERT_OK_AND_ASSIGN(
      auto input_scale,
      CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));
  TF_ASSERT_OK_AND_ASSIGN(auto input_zp, CreateRandomI32Literal({10}));
  TF_ASSERT_OK_AND_ASSIGN(
      auto output_scale,
      CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));

  function.set_name(kUncompilableFunctionName);
  const auto node_info_it =
      uncompilable_nodes.find(function.ShortDebugString());
  ASSERT_NE(uncompilable_nodes.end(), node_info_it);
  const auto& uncompilable_node_list = node_info_it->second.second;
  ASSERT_EQ(1, uncompilable_node_list.size());
  const auto& node_info = uncompilable_node_list.at(0);

  const auto& node_stack = node_info.stack_trace;

          typename Convert>
static Attribute BinaryFolder(Op *op) {
  auto lhs_op = op->getLhs().template getDefiningOp<mhlo::ConstantOp>();
  auto rhs_op = op->getRhs().template getDefiningOp<mhlo::ConstantOp>();
  if (!lhs_op || !lhs_op) return {};

  auto lhs = dyn_cast_or_null<DenseElementsAttr>(lhs_op.getValue());
  auto rhs = dyn_cast_or_null<DenseElementsAttr>(rhs_op.getValue());
  if (!lhs || !rhs) return {};

    }

    if (!tpose_arg->hasOneUse()) {
      return failure();
    }

    auto tpose_arg_type =
        llvm::dyn_cast<RankedTensorType>(tpose_arg->getResultTypes()[0]);
    auto cst_arg_type =
        llvm::dyn_cast<RankedTensorType>(cst_arg->getResultTypes()[0]);

    auto tpose_arg_rank = tpose_arg_type.getRank();
    auto cst_arg_rank = cst_arg_type.getRank();

}

TEST(BatchFunctionTest, Basic) {
  auto corert = CreateTestCoreRuntime();
  auto* host = corert->GetHostContext();
  auto [bef_buffer, bef_file] = CreateBefFile(host);
  auto* func = bef_file->GetFunction("main");
  CHECK(func);
  CHECK_EQ(func->result_types().size(), 113);
  CHECK_EQ(func->argument_types().size(), 113);

  auto arguments = CreateTestArguments(func, host);

  size_t start_map = fbb.StartMap();

  for (auto attr : attrs) {
    if (auto float_attr = mlir::dyn_cast_or_null<FloatAttr>(attr.second)) {
      fbb.Float(attr.first.data(), float_attr.getValue().convertToFloat());
    } else if (auto int_attr =
                   mlir::dyn_cast_or_null<IntegerAttr>(attr.second)) {
      fbb.Int(attr.first.data(), int_attr.getInt());
    } else if (auto bool_attr = mlir::dyn_cast_or_null<BoolAttr>(attr.second)) {

                                   SmallVector<TransposeOp, 2>* transpose_ops) {
  for (auto it = transpose_ops->begin(); it != transpose_ops->end(); ++it) {
    auto tranpose_op = *it;
    for (auto tranpose_operand : tranpose_op.getOperands()) {
      auto ranked_tranpose_type =
          mlir::dyn_cast_or_null<RankedTensorType>(tranpose_operand.getType());
      if (!ranked_tranpose_type) continue;