@ops.RegisterGradient("XlaClusterOutput")
def _XlaClusterOutputGrad(_, grad):
  del grad  # unused
  raise RuntimeError("Gradient computation of graph in xla.compile() is "
                     "prohibited because it can cause performance degradation."

      llvm::cl::desc("Choose target opset."),
      llvm::cl::values(
          clEnumValN(OpSet::TF, "TF",
                     "Uses TF ops that mimic quantization behavior"),
          clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
          clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
                     "Uses TF Uniform Quantized ops"))};
};

llvm::StringRef InsertQuantizedFunctionsPass::GetFunctionLibrary(

}

absl::Status BuildHloFromGraph(Graph& graph, bool use_output_shapes) {
  xla::XlaBuilder builder(
      ::testing::UnitTest::GetInstance()->current_test_info()->name());
  mlir::MLIRContext mlir_context;
  llvm::SmallVector<xla::XlaOp, 4> xla_params;
  std::vector<xla::XlaOp> returns(1);
  return BuildHloFromGraph(graph, builder, mlir_context, xla_params, returns,

  }

  JitCompilationListener* listener() const { return listener_; }

  // Returns a test graph that will split into two XLA clusters (due to a node
  // with _XlaCompile = false).
  GraphDef GetTestGraph(const PartialTensorShape& input_shape);

  // Runs the graph using specified batch size both with and without XLA JIT
  // compilation. Returns an error if the results between the two do not match.

// RUN: tf-opt "-xla-legalize-tf=use-tf2xla-fallback=true device-type=XLA_CPU_JIT" -verify-diagnostics %s | FileCheck --check-prefix SUPPORTED_FALLBACK_DEVICE %s
// RUN: tf-opt "-xla-legalize-tf=use-tf2xla-fallback=true" %s | FileCheck --check-prefix UNSPECIFIED_FALLBACK_DEVICE %s
// RUN: tf-opt "-xla-legalize-tf=use-tf2xla-fallback=true device-type=INVALID_DEVICE_TYPE" %s | FileCheck --check-prefix UNSUPPORTED_FALLBACK_DEVICE %s

// We run this test four times:

#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_CLIENT_H_

#include <optional>
#include <string>
#include <variant>

#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/executable_build_options.h"

namespace tensorflow {

template <typename ExecutableType, typename ClientType>
class DeviceCompilerClient {
 public:
  DeviceCompilerClient() = default;

    const std::vector<TensorShape>& arg_shapes,
    std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes);

// Compiles a serialized MLIR module into XLA HLO, generates all accompanying
// metadata and stores them in CompilationResult.
absl::StatusOr<XlaCompilationResult> LegalizeWithMlirBridge(
    const tpu::MlirToHloArgs& computation,

#include "absl/container/flat_hash_map.h"
#include "absl/numeric/bits.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/status_macros.h"
#include "xla/statusor.h"
#include "tensorflow/core/framework/types.h"

namespace tensorflow {
namespace jit {
class DeviceInfoCache;
class DeviceSet;

  NameAttrList function;
  function.set_name("foo");

  ResourceMgr resource_mgr("");

  const XlaCompiler::CompilationResult* compilation_result = nullptr;
  xla::PjRtLoadedExecutable* pjrt_executable = nullptr;
  xla::PjRtClient* pjrt_client = nullptr;

  TF_EXPECT_OK(CompileToPjRtLoadedExecutable(
      device, platform_info, function, SampleArgsForAddXY(),
      DeviceCompileMode::kStrict, /*has_ref_vars=*/true,

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

module attributes {tf.versions = {producer = 179 : i32}} {
  func.func @main(%arg0: tensor<128x8xf32> {mhlo.sharding = "bad_sharding"}) {
    func.return
  }
}