// RUN: not tf-tfrt-opt -tfrt-deduplicate-functions-invoked-by-batch-function %s 2>&1 | FileCheck %s

// This test verifies the error when two functions are different but invoked by
// the batch functions with same shared_name.

func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {

==============================================================================*/
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tfrt/analysis/tensor_array_side_effect_analysis.h"

namespace tensorflow {
namespace tfrt_compiler {
namespace {

class TestTensorArraySideEffectAnalysis
    : public mlir::PassWrapper<TestTensorArraySideEffectAnalysis,

// RUN: tf-tfrt-opt -split-input-file -tfrt-reconfig-batch-op="tfrt-min-num-batch-threads=2 tfrt-min-max-enqueued-batches=3" %s | FileCheck %s --dump-input=always

// -----

// The num_batch_threads is lowered bound to 2 from the original attribute of 1

// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>

// RUN: tf-tfrt-opt -tfrt-lower-cluster-to-runtime-ops-tpu -split-input-file -verify-diagnostics %s | FileCheck %s

module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {

  // CHECK-LABEL: @converts_cluster
  func.func @converts_cluster() {
    // CHECK: %0:2 = "tf_device.launch"() <{{.*}}> ({

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"

#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project

  return tensorflow::io::JoinPath(tensorflow::testing::TensorFlowSrcRoot(),
                                  kTestData, saved_model_dir);
}

// This value parameterized test allows us to test both TFRT
// and non TFRT runtimes.
// https://github.com/google/googletest/blob/dcc92d0ab6c4ce022162a23566d44f673251eee4/googletest/docs/advanced.md#value-parameterized-tests
class CSavedModelAPITest : public ::testing::TestWithParam<bool> {};

#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_UTIL_H_

#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project

namespace tfrt {
namespace fallback_async {

bool IsArgConsumedByFallback(mlir::func::FuncOp func, int arg_index);

void ForEachArgConsumedByFallback(
    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action);

// RUN: tf-tfrt-opt -tf-to-tfrt %s | FileCheck %s --dump-input=fail

// CHECK-LABEL: func @derived_attrs
func.func @derived_attrs(
  %serialized: tensor<?x!tf_type.string>,
  %names: tensor<0x!tf_type.string>,
  %sparse_keys: tensor<0x!tf_type.string>,
  %dense_keys: tensor<1x!tf_type.string>,
  %ragged_keys: tensor<0x!tf_type.string>,
  %dense_default: tensor<0xi64>) -> tensor<?xi64> {

  %dense_value =

#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tfrt/analysis/cost_analysis.h"

namespace tensorflow {
namespace tfrt_compiler {

class TestCostAnalysis
    : public mlir::PassWrapper<TestCostAnalysis,
                               mlir::OperationPass<mlir::func::FuncOp>> {
  llvm::StringRef getArgument() const final {
    return "tfrt-test-cost-analysis";
  }
  llvm::StringRef getDescription() const final {

#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tensorflow/core/tfrt/common/global_state.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tensorflow/core/tpu/tpu_defs.h"
#include "tsl/framework/device_type.h"

namespace tensorflow {
namespace {