// -----

// ============================================================================
// The following functions tests example quantization patterns outputted from
// StableHLO Quantizer. These patterns should be legalized early directly
// to fused tflite ops.
// ============================================================================

// Tests that a simple per-channel quantized `stablehlo.dot_general` is properly

// Create a pass that legalizes MHLO to TF dialect.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeHloToTfPass();

// Creates a pass which replaces a splat constant tensor with a BroadcastInDim
// op.
std::unique_ptr<OperationPass<ModuleOp>> CreateUnfoldSplatConstantPass();

// Create a pass that legalizes MHLO to TFLite dialect.

    pass_manager->addPass(mlir::TF::CreateTFShapeInferencePass());
  }

  // Legalize while early to allow further constant folding.
  // TODO(jpienaar): This may not actually matter as we do canonicalization
  // after the legalize below, for now it needs to be below the above passes
  // that work on TF dialect and before inliner so that the function calls in

A. Unique TensorFlower <******@****.***> 1713119208 -0700

namespace mlir {
namespace TF {

//===----------------------------------------------------------------------===//
// The pass to legalize the quantization emulation ops from TF.
//
namespace {

// Legalize TF quantization emulation ops to that in Quant ops dialect.
struct LegalizeTFToQuant
    : public PassWrapper<LegalizeTFToQuant, OperationPass<func::FuncOp>> {

// RUN: tf-opt -tfl-legalize-random %s | FileCheck %s


// CHECK-LABEL:   func @tfl_wrapped_jax_random_normal(
// CHECK-SAME:                                        %[[RNG:.*]]: tensor<2xui32>) -> tuple<tensor<3x4xf32>> {
// CHECK:           %[[VAL_0:.*]] = stablehlo.constant dense<[3, 4]> : tensor<2xi32>

A. Unique TensorFlower <******@****.***> 1712604801 -0700

auto* mlir_failed_legalization_op_count =
    tensorflow::monitoring::Counter<1>::New(
        "/tensorflow/core/tf2xla/"
        "mlir_second_phase_failed_legalization_op_count",
        "Counts which op fails to legalize", "op_name");

auto* mlir_non_static_op_count = tensorflow::monitoring::Counter<1>::New(
    "/tensorflow/core/tf2xla/"
    "mlir_second_phase_non_static_op_count",

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

namespace tensorflow {
namespace tf2xla {
namespace v2 {

// Legalizes the given mlir::Module into XLA HLO. If successful, returns the
// compiled XLA HLO. V1 of the tf2xla uses MLIR whereas V0 does not use MLIR.
//
// Inputs:
//  computation - The MLIR module op. It currently takes in

#include "xla/client/compile_only_client.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Legalize the given MLIR module to XLA HLO using a combination of the MLIR
// Bridge and XlaBuilder
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
    const tpu::MlirToHloArgs& computation,