return value.getDefiningOp();
}

// Check if the TPUCopyWithDynamicShapeOp is valid.
// 1. The op should be wrapped inside a launch op.
// 2. The wrapped launch op should be placed on CPU.
LogicalResult CheckOpIsValid(Operation* op) {
  auto launch_op = llvm::dyn_cast<tf_device::LaunchOp>(op->getParentOp());
  if (!launch_op) {
    op->emitError() << "TPUCopyWithDynamicShapeOp is not in a launch";
  }

  solver.load<dataflow::SparseConstantPropagation>();
  solver.load<DeviceDataflowAnalysis>(symbolTables);
  if (failed(solver.initializeAndRun(module))) return signalPassFailure();

  module->walk([&](TF::TPUCopyWithDynamicShapeOp op) {
    const Device *state;
    for (auto result : op->getResults()) {
      state = solver.lookupState<Device>(result);
      if (state) break;
    }
    if (!state || !state->hasDevice()) {

// partitioned resource variables.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateTPUResourceReadsWritesPartitioningPass();

// Creates a pass that looks for usage of the result of
// TPUCopyWithDynamicShapeOp and annotate these values to be dynamic shape. This
// ensures that the generated tpu program has the correct inputs annotation.
std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUAnnotateDynamicShapeInputsPass();

def ExtractTPUCopyWithDynamicShapeOpPass : Pass<"tf-extract-tpu-copy-with-dynamic-shape-op", "mlir::func::FuncOp"> {
  let summary = "Extract the TPUCopyWithDynamicShapeOp out of the host launch and place it on device launch";

  let description = [{
    This pass looks for TPUCopyWithDynamicShapeOp which wraps in a
    `tf_device.launch` with host device attribute. It extracts the ops and wrap

  }) : () -> tensor<f32>
  return %0 : tensor<f32>
}
```
### `-tf-extract-tpu-copy-with-dynamic-shape-op`

_Extract the TPUCopyWithDynamicShapeOp out of the host launch and place it on device launch_

This pass looks for TPUCopyWithDynamicShapeOp which wraps in a
`tf_device.launch` with host device attribute. It extracts the ops and wrap
them in `tf_device.launch` with tpu device attribute so that ops can be

}

def TF_TPUAnnotateTensorsWithDynamicShapeOp : TF_Op<"TPUAnnotateTensorsWithDynamicShape", [Pure]> {
  let summary = [{
Placeholder op which takes the output of TPUCopyWithDynamicShapeOp and pass
them to the following tpu ops.
  }];

  let description = [{
This op serves as an annotation for the dynamic shaped tensor and will be
removed during the bridge rewrite.
  }];