// (a) read-only. These reads can be replaced by a hoisted read created
  //        before the WhileOp (similar to if and case).
  // (b) written: since the value is written in the loop (which can only in
  //        loop body, all these will become loop variables. Since all resource
  //        variables are removed from the loop variabled during
  //        canonicalizationW, we need to create new operand/result slots. The

// Test to ensure that a registered graph optimization pass is only executed
// once (i.e., on the main function side) in running distributed functions.
// This test creates a cluster with two workers, create a variable on the
// second worker, and run a distributed function (VariableAddFunction) whose ops
// span the local and remote workers. If the graph optimization pass is executed

  if (only_one_return_value) {
    return_types.resize(1);
  }

  auto type = FunctionType::get(region.getContext(), input_types, return_types);

  // Create new function and extract region body into the function.
  auto outlined_func = builder.create<func::FuncOp>(loc, name, type);
  Region& func_region = outlined_func.getBody();
  func_region.takeBody(region);
  Block& first_block = func_region.front();

  auto graph_op = builder.create<tf_executor::GraphOp>(
      loc, func.getFunctionType().getResults());
  graph_op.getBody().push_back(new Block);
  builder.setInsertionPointToEnd(&graph_op.GetBody());
  auto island = builder.create<tf_executor::IslandOp>(
      loc, func.getFunctionType().getResults(),
      tf_executor::ControlType::get(&getContext()), ArrayRef<Value>());
  // Create Fetch.

    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.

    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),

  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>

  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(),

  context.loadDialect<TensorFlowDialect>();
  OpBuilder opBuilder(&context);
  OperationState state(UnknownLoc::get(opBuilder.getContext()),
                       /*OperationName=*/"tf.Const");
  mlir::Operation* op = Operation::create(state);

  bool result = IsInBridgeAcceptableDialects(op);

  EXPECT_TRUE(result);
  op->destroy();
}

TEST_F(SharedUtilsTest, IsInFunctionalDialectFails) {
  MLIRContext context;