// 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>

        source.contains("foo")
    }

    def "an element added as both a provider and a realized value is not duplicated"() {
        when:
        source.add("foo")
        source.addPending(provider("foo"))

        then:
        source.iterator().collect() == ["foo"]
    }

    def "an element added as both a provider and a realized value has the correct order"() {
        when:
        source.addPending(provider("foo"))

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 "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,

#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

  }

  if (failed(mlir::tf_saved_model::FreezeVariables(
          module_op, saved_model_bundle->GetSession()))) {
    return absl::InternalError("Failed to freeze variables.");
  }

  // Run legalize TF to StableHLO pass to convert `tf.Const` and
  // `tf.Const`->`tf.Cast` patterns after variable freezing. The TF shape
  // inference pass is also required to resolve unknown shapes in the TF dialect

    registry.insert<shape::ShapeDialect>();
  }

 public:
  StringRef getArgument() const final { return "tf-mhlo"; }
  StringRef getDescription() const final {
    return "This pass will legalize TF Ops to MHLO Ops.";
  }

 protected:
  Option<bool> skip_quantization_ops_{
      *this, "skip-quantization-ops",
      ::llvm::cl::desc("Skip quantization ops")};

  Option<bool> skip_resize_{

// RUN: odml-to-stablehlo-opt %s -tfl-legalize-hlo -split-input-file | FileCheck %s --dump-input=fail

// CHECK-LABEL: mhlo_custom_call_test__legalize_string_backend_config
func.func @mhlo_custom_call_test__legalize_string_backend_config(%arg0: tensor<1x4xf32>) -> tensor<1x8xf32> {
  %0 = mhlo.custom_call @custom_call.my_custom_op(%arg0) {
    api_version = 1 : i32,
    backend_config = "this_is_a_test_string"
  } : (tensor<1x4xf32>) -> (tensor<1x8xf32>)

          for (auto arg : block.getArguments()) {
            ConvertAndWrapUsesInUnrealizedCast(arg, converter, rewriter);
          }
        }
      }
    });
  });
  return success();
}

// Legalize StableHLO portion of program to VHLO, leaves TFL untouched
LogicalResult ApplyStablehloToVhloPatterns(ModuleOp module,
                                           bool is_func_legal) {
  MLIRContext *context = module.getContext();