// This pattern converts TensorFlow axis format to HLO axis format which
// doesn't wrap around like TensorFlow and is always positive. For this
// conversion, use the first input to get inputs rank. Other inputs need not be
// ranked.
// Defining op for `axis` is TensorFlow constant op in the pattern as during
// the conversion, original Concat op operands still refers to the old ops even

std::unique_ptr<OperationPass<ModuleOp>> CreateCompositeLoweringPass();

// Adds the HLO to TF rewrite patterns to the specified pattern list.
void PopulateLegalizeHloToTfPatterns(RewritePatternSet* patterns,
                                     MLIRContext* context);

// Adds the HLO to TFLite rewrite patterns to the specified pattern list.
void PopulateLegalizeHloToTFLitePatterns(RewritePatternSet* patterns,

};

// Converts tf.RandomShuffle op into a series of XLA HLO ops.
//
// tf.RandomShuffle shuffles tensors along the first dimension. If the input
// tensor's rank is 1, then it is translated into HLO sort op(s) according to
// indices randomly generated via HLO rng_uniform ops. Otherwise, it is
// translated into an HLO while op to first emulate shuffling indices using
// HLO dynamic_slice and dynamic_update_slice ops, then finally HLO gather

config.name = 'MLIR ' + os.path.basename(config.mlir_test_dir)

config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)

# suffixes: A list of file extensions to treat as test files.
config.suffixes = ['.cc', '.hlo', '.json', '.mlir', '.pbtxt', '.py']

# test_source_root: The root path where tests are located.
config.test_source_root = config.mlir_test_dir

# test_exec_root: The root path where tests should be run.

namespace internal {

// Compiles a serialized MLIR module and returns a serialized MLIR module of the
// result of running all the MLIR Bridge passes. If compile_to_xla_hlo is true
// then those passes include all the Legalization to XLA HLO which is returned
// in the compilation_result.
absl::StatusOr<std::string> CompileFromMlirToXlaHlo(
    bool lower_to_xla_hlo, const tpu::MlirToHloArgs& computation,

// RUN: tf_tfl_translate --enable-hlo-to-tf-conversion --input-mlir %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s --check-prefix=CHECK-ROUNDTRIP


module {
  func.func public @main( %arg0: tensor<i64>) ->  tensor<i64> {
    %0 = func.call @test_add_roundtrip(%arg0) : (tensor<i64>) -> tensor<i64>

    return %0 : tensor<i64>
  }


  // CHECK-LABEL: func.func private @test_add_roundtrip

  let options = [
    Option<"default_patterns_", "default-patterns", "bool", /*default=*/"true",
           "Activate the default TF->TF patterns">,
    Option<"pre_hlo_patterns_", "pre-hlo-patterns", "bool", /*default=*/"true",
           "Activate the pre-HLO TF->TF patterns">,
  ];
}

def TestResourceAliasAnalysis : Pass<"tf-test-resource-alias-analysis", "mlir::func::FuncOp"> {

  let description = [{
    Gathers metrics about the input MLIR to Phase 2 of the TFXLA Bridge, which
    does a strict semantic lowering from Tensorflow ops to XLA HLO.
  }];

  let constructor = "tensorflow::tf2xla::internal::CreateInputLoweringMetricsPass()";

  // Whether to enable the hlo/stablehlo to tf conversion. This also supports
  // the case where a saved model contains both TF module and serialized
  // StableHLO module.
  bool enable_hlo_to_tf_conversion = false;
  // Whether to disable the direct hlo/stablehlo to Tensorflow Lite conversion.
  //
  // This prevents from directly converting from HLO to TFLite without going

    std::vector<ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes) {
  LOG_FIRST_N(INFO, 1)
      << "Compiling MLIR computation to XLA HLO using MLIR tf2xla bridge in "
         "the op by op fallback mode. This is Phase 2 of the TF2XLA Bridge. "
         "Old (non-MLIR) bridge may be used in case of unsupported feature "