#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

// Pass that co-locates resource ops that use composite device resources
// (packed tensors) with the underlying physical TPU device.
struct TPUColocateCompositeResourceOps
    : public impl::TPUColocateCompositeResourceOpsPassBase<
          TPUColocateCompositeResourceOps> {
  void runOnOperation() override;
};

  auto status_or_device_coodinates =
      tensorflow::GetDeviceCoordinates(device_assignment_attr);
  if (!status_or_device_coodinates.ok())
    return cluster.emitError()
           << "error in fetching tpu device coordinates: "
           << status_or_device_coodinates.status().message();

  status_or_tpu_device_assignment =
      tensorflow::GetTPUCompilationAndExecutionDevices(

#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/tpu/tpu_defs.h"

namespace tensorflow {
namespace {

TEST_F(OpsTestBase, CreateSingleOpGraph) {
  TF_EXPECT_OK(NodeDefBuilder("identity_op", "Identity")
                   .Input(FakeInput(DT_FLOAT))

                                 RewritePatternSet* patterns) {
  for (auto op : context->getRegisteredOperations())
    op.getCanonicalizationPatterns(*patterns, context);
}

// Adds the list of ops that are supported on TPU through constant folding which
// may depend on the inputs shapes not known at this point. Such ops may not
// have any legalization or canonicalization patterns but shouldn't be marked
// for outside compilation.
//

  // TODO: b/230572023 - Consider improving shape inference for While op instead
  // of dropping the attribute. This need not be correct for models not trained
  // on TPU.

  // Optimizes TF graph via cleanups, merges, rewrites, constant folding,
  // and edge case handling where possible.
  pm.addNestedPass<func::FuncOp>(TF::CreateDropWhileShapeInvariantPass());

// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-tpu-update-embedding-enqueue-op-inputs | FileCheck %s

// CHECK-LABEL: func @check_enqueue_ops_update_for_eval
// CHECK-SAME: %[[ARG_0:[a-z0-9]*]]: tensor<?x2xi32>
// CHECK-SAME: %[[ARG_1:[a-z0-9]*]]: tensor<?x2xi32>
// CHECK-SAME: %[[ARG_2:[a-z0-9]*]]: tensor<?x2xi32>
// CHECK-SAME: %[[ARG_3:[a-z0-9]*]]: tensor<?xi32>
// CHECK-SAME: %[[ARG_4:[a-z0-9]*]]: tensor<?xi32>

// RUN: tf-opt -split-input-file -verify-diagnostics -tf-extract-tpu-copy-with-dynamic-shape-op %s | FileCheck %s

// Test that extract TPUCopyWithDynamicShape from host launch to device launch.

module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:COMPOSITE:0", "/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0"}} {

  quantization_config.mutable_static_range_ptq_preset()
      ->set_enable_per_channel_quantized_weight(
          quantization_options.enable_per_channel_quantization());
  // When targeting server TPUs quantized types should be unpacked into
  // integer ops.
  quantization_config.mutable_pipeline_config()->set_unpack_quantized_types(
      true);
  *quantization_config.mutable_debugger_config() =

  // Get the device context given the index.
  absl::StatusOr<DeviceContext*> GetDeviceContextWithIndex(int index);

  // Instructs this XlaDevice to set a AcceleratorDeviceInfo, which holds extra
  // information for GPU and TPU devices.
  Status UseAcceleratorDeviceInfo() TF_LOCKS_EXCLUDED(mu_);

  // Instructs this XlaDevice to return 'sync_on_completion' for
  // AllowsSyncOnCompletion().

#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "mlir/Support/Timing.h"  // from @llvm-project

namespace tensorflow {

// Logger for logging MLIR modules before and after passes in MLIR TPU bridge.
//
// The IR logging can be restricted to a particular set of pass invocations via
// filters that are specified with the `MLIR_BRIDGE_LOG_PASS_FILTER` and
// `MLIR_BRIDGE_LOG_STRING_FILTER` environment variables.