llvm::SmallVector<Value, 2> new_operands;
    FilterOutBlockArgControlDep(op->getOperands().drop_front(), new_operands);

    auto source_op = builder.create<tf_executor::NextIterationSourceOp>(
        op->getLoc(), op->getOperand(0).getType());
    builder.create<tf_executor::NextIterationSinkOp>(
        op->getLoc(), source_op.getToken(), /*input=*/op->getOperand(0),
        /*controlInputs=*/new_operands);

  // Create SessionInitializerOp; should reference main function.
  StringAttr func_name = main_func.getSymNameAttr();
  llvm::SmallVector<mlir::Attribute, 2> func_names = {
      mlir::SymbolRefAttr::get(builder.getContext(), func_name)};
  builder.create<tf_saved_model::SessionInitializerOp>(
      module->getLoc(), builder.getArrayAttr(func_names));
  // Create AssetOp; this holds the checkpoint_path.

// control flow op into an i1 value.
static Value LowerCondition(Location loc, Value value, OpBuilder* builder) {
  Value zero_d = builder->create<ToBoolOp>(loc, value);
  Value scalar = builder->create<tensor::ExtractOp>(loc, zero_d);
  return scalar;
}

// Calls the function `fn` with arguments provided by the given function and

    indices[i] = dimension_numbers.start_index_map()[i];
  }

  // Fill elements from start_indices with start_index_map
  Value scattered_start_indices = builder.create<TF::TensorScatterUpdateOp>(
      loc, empty_start_indices,
      /*indices=*/
      builder.create<TF::ReshapeOp>(
          loc, RankedTensorType::get({index_map_size, 1}, builder.getI64Type()),
          Create1DConstValue<int64_t>(builder, loc, indices),

    Value new_reshape_shape_value = rewriter.create<arith::ConstantOp>(
        inner_reshape_op->getLoc(),
        GetI32ElementsAttr(new_reshape_shape_i32, &rewriter));

    auto new_inner_reshape_op = rewriter.create<TFL::ReshapeOp>(
        inner_reshape_op->getLoc(),
        inner_reshape_input, new_reshape_shape_value);

    // Create a new reshape_op to replace the old inner reshape_op.

  Type resultType = RankedTensorType::get(shape, element_type);
  auto constant_attr = DenseElementsAttr::get(shape_spec_type, shape);
  auto shape_tensor =
      rewriter.create<TF::ConstOp>(loc, shape_spec_type, constant_attr);
  return rewriter.create<TF::ReshapeOp>(loc, resultType, /*tensor=*/value,
                                        /*shape=*/shape_tensor);
}

template <typename BatchMatMulOpType>

      if (!iv.use_empty()) {
        // iv' = iv + step * i;
        Value iter = rewriter.create<arith::ConstantIndexOp>(loc, i);
        Value step_cst =
            rewriter.create<arith::ConstantIndexOp>(loc, step.getSExtValue());
        Value stride = rewriter.create<arith::MulIOp>(loc, step_cst, iter);
        Value iv_unroll =
            rewriter.create<arith::AddIOp>(loc, mapping.lookup(iv), stride);
        mapping.map(iv, iv_unroll);
      }

    // Creates a _Retval node for the src node of edge, and add it to results_,
    // if none exists yet. If a new _Retval node is created, also adds the edge
    // within the subgraph from the src to the _Retval node.
    Status RecordResult(
        const Edge* edge,
        const absl::flat_hash_map<const Node*, Node*>& node_images);

  const OwningOpRef<mlir::stablehlo::ConstantOp> constant_op =
      builder_.create<mlir::stablehlo::ConstantOp>(
          builder_.getUnknownLoc(), builder_.getI32IntegerAttr(0));
  EXPECT_TRUE(IsStablehloOp(*constant_op));
}

TEST_F(StablehloTypeUtilsTest, IsStablehloOpFailsWithArithOp) {
  const OwningOpRef<mlir::arith::ConstantOp> constant_op =
      builder_.create<mlir::arith::ConstantOp>(builder_.getUnknownLoc(),

  }

  // Create the function.
  MLIRContext* context = ops[0]->getContext();
  StringRef dialect = ops[0]->getName().getDialectNamespace();
  OpBuilder builder(context);
  // Every ModuleOp has at least one region and one block.
  Block* first_block = &module->getRegion(0).front();
  builder.setInsertionPointToEnd(first_block);
  auto func = builder.create<mlir::func::FuncOp>(