default_visibility = [
        "//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
        "//tensorflow/compiler/mlir:__pkg__",
        "//tensorflow/compiler/mlir/tensorflow/transforms:__pkg__",
        "//tensorflow/compiler/mlir/tf2xla/api:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/compiler/tf2xla:__pkg__",

#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tsl/framework/allocator.h"
#include "tsl/lib/core/status_test_util.h"
#include "tsl/platform/status.h"
#include "tsl/platform/statusor.h"

#ifdef TF_MLRT_OPS
#else
#define TF_MLRT_OPS

include "mlir/Interfaces/SideEffectInterfaces.td"
include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td"
include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"

def CreateOp: TensorflowMlrt_Op<"createop", []> {
  let summary = "The Fallback CreateOp";

  let description = [{

#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/framework/device_factory.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/tfrt/common/async_value_tensor.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tsl/framework/device_id_utils.h"

namespace tensorflow {
namespace {

absl::StatusOr<std::unique_ptr<xla::PjRtBuffer>> HostTensorToPjRtBuffer(

    Status s = tensorflow::StatusFromTF_Status(status.get());
    CHECK_EQ(tensorflow::errors::OK, s.code()) << s.message();

    // Set the runtime impl, Core RT vs TFRT.
    AbstractContext* ctx_raw = nullptr;
    s = BuildImmediateExecutionContext(UseTfrt(), &ctx_raw);
    CHECK_EQ(tensorflow::errors::OK, s.code()) << s.message();
    ctx_.reset(ctx_raw);
  }
};

// RUN: tf-tfrt-opt -split-input-file -tfrt-optimize-tf-control-flow-side-effect %s | FileCheck %s

func.func @no_side_effect_cond(%arg: tensor<i32>) -> tensor<i1> {
  %0 = "tf.Const"() {value = dense<16> : tensor<i32> } : () -> tensor<i32>
  %1 = "tf.Less"(%arg, %0) : (tensor<i32>, tensor<i32>) -> tensor<i1>
  func.return %1 : tensor<i1>
}

func.func @no_side_effect_body(%arg: tensor<i32>) -> tensor<i32> {

// RUN: tf-tfrt-opt --split-input-file -pass-pipeline='builtin.module(pre-parallel-tf-to-mlrt{use-tpu-host-allocator-for-inputs=true},tf-mlrt-parallelization{tfrt-cost-threshold=4},tf-to-mlrt)'  %s | FileCheck %s --dump-input=fail --dump-input-filter=all

func.func @callee(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
  func.return %arg0: tensor<i32>
}

// CHECK-LABEL: func @batch_function
func.func @batch_function(%arg0: tensor<i32>) -> (tensor<i32>) {

// RUN: tf-tfrt-opt -tfrt-test-cost-analysis -verify-diagnostics %s | FileCheck %s

// CHECK-LABEL: test_cheap_ops_0
func.func @test_cheap_ops_0(%arg: tensor<?x!tf_type.string>) -> (tensor<?x8xf32>) {
    // expected-remark@+1 {{Cost: 1}}
    %0 = "tf.Const"() {value = dense<> : tensor<0xi64>} : () -> tensor<0xi64>
    // expected-remark@+1 {{Cost: 1}}

  cat > $BATS_TEST_TMPDIR/ignore_deps_for_these_packages <<EOF
//tensorflow/lite
//tensorflow/compiler/mlir/lite
//tensorflow/compiler/mlir/tfrt
//tensorflow/core/runtime_fallback
//tensorflow/core/tfrt
//tensorflow/python/kernel_tests/signal
//tensorflow/examples
//tensorflow/tools/android
//tensorflow/python/eager/benchmarks
EOF

    ::testing::Combine(::testing::Values("graphdef", "mlir"),
                       /*tfrt*/ ::testing::Values(false),
                       /*use_function*/ ::testing::Values(true, false)));
#else
INSTANTIATE_TEST_SUITE_P(
    UnifiedCAPI, CppGradients,
    ::testing::Combine(::testing::Values("graphdef", "mlir"),
                       /*tfrt*/ ::testing::Values(false),