const Type expressed_type = value.getType();
    // The value needs to be requantized. A Quantize op will be created to use
    // it as the operand and replace its uses.
    const Type new_type = state.params.castFromExpressedType(expressed_type);
    if (!new_type) return;
    auto requantize_op =
        builder_.create<quantfork::QuantizeCastOp>(loc, new_type, value);
    value.replaceAllUsesWith(requantize_op);

  if (sub_op) {
    sub_op->replaceAllUsesWith(island.getOutputs());
    sub_op->moveBefore(block, block->begin());
    island_builder.create<tf_executor::YieldOp>(loc, sub_op->getResults());
  } else {
    island_builder.create<TF::NoOp>(island.getLoc(), TypeRange{}, ValueRange{});
    island_builder.create<tf_executor::YieldOp>(loc, ValueRange{});
  }
  return island;
}

  // Guarantee all functions have one use, which enables more exact shape
  // inference.
  pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
  // Run shape inference so that tf_executor/tf_device ops created later will
  // likely to inherit more concrete types.
  pm.addPass(mlir::TF::CreateTFShapeInferencePass());
  pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUPartitionedOpConversionPass());
  pm.addNestedPass<FuncOp>(

    return tokens[0];
  } else {
    return builder.create<AfterAllOp>(loc, original_token.getType(), tokens)
        .getResult();
  }
}

// Replaces `tf._XlaHostComputeMlir` with individual `mhlo.send` and `mhlo.recv`
// ops per operand and result. Unique Channel IDs are assigned per transfer.
// Sink tokens are created across all `mhlo.send` ops first and then by
// all `mhlo.recv` ops.

      if (pjrt_tensor_buffer != nullptr) {
        args->push_back(pjrt_tensor_buffer->pjrt_buffer());
      } else {
        // Creates a PjRtBuffer from DeviceMemoryBase. The newly created
        // PjRtBuffer needs to be persisted till XLA execution is completed.
        auto dmem = se::DeviceMemoryBase(
            const_cast<char*>(tensor->tensor_data().data()),

}

// Wraps block in a Launch. External uses of ops in the block will be return
// values of the Launch and remapped to the Launch results. If `before` is set
// to true, the Launch is created before `op`. Otherwise the Launch is created
// after `op`.
mlir::tf_device::LaunchOp CreateLaunchForBlock(OpBuilder* builder,
                                               Operation* op, bool before,

                 Out(NodeWith(
                     Op("Slice"),
                     Name("begin/static_shaped_slice/static_shaped_slice"))),
                 _)));
}

// New constants being created need to have control dependencies copied to
// ensure correct control flow analysis in TF V2.
TEST(SliceToDynamicSliceRewriteTest, WithControlDepsToConstant) {
  Scope root = Scope::NewRootScope()

  void MarkAsResult(AbstractTensorHandle* gradient) const override;

  void DeleteGradient(AbstractTensorHandle* gradient) const override;

 private:
  // The context where the aggregation op `Add` is to be created.
  AbstractContext* ctx_;
};

// Returns the number of elements in the gradient tensor.
int64_t TapeVSpace::NumElements(AbstractTensorHandle* tensor) const {

  }
  void runOnOperation() override;
};

// Creates a call to function `func` in region `caller_region`. Use `args` as
// the call arguments, and terminate the region with a yield. The arguments are
// cast to the required type before the call. `use_region_args` control whether
// the input arguments are used as is (for IfOp) or block arguments of the same
// type as the input arguments are created and then used as call arguments (for
// While).

      builder.clone(*op, mapping);
    }
    // Creates the ReturnOp so that the per-host function returns the
    // correct values of the cloned operations.
    llvm::SmallVector<Value, 4> results_after_mapping;
    for (Value result : metadata.results) {
      results_after_mapping.push_back(mapping.lookupOrDefault(result));
    }
    builder.create<func::ReturnOp>(loc, results_after_mapping);