// 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

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

    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>

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

  func->setAttr(kTFImplements, attr);
  OpBuilder builder(func.getBody());
  std::string empty_option_buffer;
  auto op = builder.create<CustomOp>(
      func.getLoc(), func.getFunctionType().getResults(), func.getArguments(),
      api, CustomOption(&builder, empty_option_buffer));
  builder.create<func::ReturnOp>(func.getLoc(), op.getResults());
  return success();
}

LogicalResult VerifyNgrams(func::FuncOp func) {

namespace stablehlo::quantization::io {

absl::StatusOr<std::string> GetLocalTmpFileName(tsl::Env* const env) {
  std::string tmp_fname{};
  if (!env->LocalTempFilename(&tmp_fname)) {
    return absl::InternalError("Failed to create tmp file name.");
  }

  return tmp_fname;
}

absl::StatusOr<std::string> GetLocalTmpFileName() {
  return GetLocalTmpFileName(tsl::Env::Default());
}