// RUN: tf-mlir-translate -mlir-tf-to-hlo-text %s -tf-input-shapes=128,10:10,1024:128,1024 -tf-xla-emit-use-tuple-args -tf-xla-emit-return-tuple | FileCheck %s

module attributes {tf.versions = {producer = 179 : i32}} {
  func.func @main(%arg0: tensor<128x10xf32> {mhlo.sharding = "\08\03\1A\02\01\02\22\02\00\01"}, %arg1: tensor<10x1024xf32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg2: tensor<128x1024xf32> {mhlo.sharding = ""}) {
    func.return
  }
}

#include "tensorflow/core/tpu/kernels/tpu_compile.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"

namespace tensorflow {
namespace tf2xla {
namespace v1 {

// Compiles the given Tensorflow graph into xla::HLO. The result is in
// compilation_result. If the input computation is in MLIR, it will be
// converted to a Tensorflow graph. Otherwise, the graph compiler will be run.
absl::Status CompileTensorflowGraphToHlo(

  auto is_tfrt_tpu_supported_stage = [](IrExportStage stage) {
    return stage == IrExportStage::HLO ||
           stage == IrExportStage::HLO_NO_METADATA ||
           stage == IrExportStage::HLO_SERIALIZED;
  };
  // TODO(b/238830423): support GetCompilerIr on TFRT TPU device for stages
  // that requires compilation from HLO to executable.
  if (device_type != DEVICE_CPU && stream == nullptr &&

// RUN: tf-mlir-translate -mlir-tf-to-hlo-text %s -tf-input-shapes=: -tf-xla-emit-return-tuple | FileCheck %s


// TODO(b/259459405): Remove this test along with the upstream refactoring to
// avoid non TF inputs.
// This is not a supported mode.
module attributes {tf.versions = {producer = 179 : i32}} {
  func.func @main(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>

// RUN: tf-opt -tf-to-hlo-pipeline %s | FileCheck %s

// Verifies that constants generated post shape inference are propagated.
// get_shape result in this test.
module attributes {tf.versions = {producer = 179 : i32}} {

  // CHECK-LABEL: func @main
  func.func @main(%arg0: tensor<10x19xf32>, %arg1: tensor<19x10xf32> {mhlo.is_same_data_across_replicas = true}) -> tensor<?xi64> {
    %0 = "tf.Shape"(%arg0) : (tensor<10x19xf32>) -> tensor<2xi64>

include "mlir/Pass/PassBase.td"

def LegalizeTFCommunicationPass : Pass<"xla-legalize-tf-communication", "ModuleOp"> {
  let summary = "Legalize TF/XLA communication ops (TensorFlow dialect) to the HLO "
           "dialect";
  let dependentDialects = ["mhlo::MhloDialect", "sparse_tensor::SparseTensorDialect"];
  let constructor = "::mlir::mhlo::CreateLegalizeTFCommunicationPass()";
  let description = [{

// RUN: tf-mlir-translate -mlir-tf-to-hlo-text %s -tf-input-shapes=: | FileCheck %s

module attributes {tf.versions = {producer = 179 : i32}} {
  func.func @main() -> (tensor<300x?xf32>, tensor<300x?xf32>) {
    %elem_shape = "tf.Const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
    %size = "tf.Const"() {value = dense<300> : tensor<i32>} : () -> tensor<i32>

// RUN: tf-mlir-translate -mlir-tf-graph-to-hlo-text %s -tf-input-shapes=2:2 -tf-input-data-types=DT_FLOAT,DT_FLOAT -tf-xla-input-types=parameter,resource -tf-xla-emit-return-tuple | FileCheck %s

module attributes {tf.versions = {producer = 511 : i32}} {
  func.func @main(%arg0: tensor<*xf32>, %arg1: tensor<*x!tf_type.resource>) {
    tf_executor.graph {

      tensorflow::testing::TmpDir(), "my_test_prefix"));

  // Run again and these should all hit in the persistent cache despite having
  // altered the persistent cache entries' HLO modules (disabled strict
  // signature checks).
  listener()->ClearListenerHistory();
  for (int b = 1; b < 4; ++b) {
    TF_ASSERT_OK(ExecuteWithBatch(graph, b));
  }

# MLIR dialects

## Overview


To separate different hardware and software targets, MLIR has “dialects”,
including:

* TensorFlow IR, which represents all things possible in TensorFlow graphs.
* XLA HLO IR, which is designed to take advantage of XLA’s compilation
  abilities (with output to, among other things, TPUs).
* An experimental affine dialect, which focuses on
  [polyhedral representations](https://en.wikipedia.org/wiki/Polytope_model)