// Update input type.
  auto transform_result_type =
      RankedTensorType::get(new_shape, getElementTypeOrSelf(cast_input));
  cast_input.setType(transform_result_type);
  auto block_arg = mlir::dyn_cast<mlir::BlockArgument>(cast_input);
  auto cast_op_input = dyn_cast_or_null<TF::CastOp>(cast_input.getDefiningOp());
  while (block_arg || cast_op_input) {
    if (block_arg) {
      // Change on device function type/shape.

          while_op, while_op.getBodyAttr());
      for (auto &arg : while_op->getOpOperands()) {
        BlockArgument cond_arg = cond.getArgument(arg.getOperandNumber());
        this->join(getLatticeElement(cond_arg), *getLatticeElement(arg.get()));
        BlockArgument body_arg = body.getArgument(arg.getOperandNumber());
        this->join(getLatticeElement(body_arg), *getLatticeElement(arg.get()));
      }

// `tf.ReadVariableOp` and `tf.AssignVariableOp` only.
mlir::LogicalResult ValidateResourceArgument(func::FuncOp function,
                                             BlockArgument resource_arg,
                                             TF::ResourceType resource_type) {
  if (resource_type.getSubtypes().size() != 1)
    return function.emitError()
           << "expects resource type of argument "

    return false;
  llvm::SmallSetVector<Value, 8> aliases =
      resource_alias_analysis.GetResourceAliases(resource_iterator);

  auto is_generator = [](Value val) {
    if (mlir::isa<BlockArgument>(val)) return true;
    Operation* definition = val.getDefiningOp();
    return definition->getNumOperands() == 0 &&
           definition->getNumResults() == 1;
  };

    }
  }

  // Collect all those freezable. This is an extra scan but allows for the
  // partial behavior from `allow_mutable_tensor`.
  DenseMap<BlockArgument, bool> freezeable;
  for (auto func : module.getOps<func::FuncOp>()) {
    for (BlockArgument val : func.getArguments()) {
      if (!mlir::isa<TF::ResourceType>(getElementTypeOrSelf(val.getType())))
        continue;

    // defining op.
    if (def_op != nullptr &&
        (isa<TF::TPUExecuteOp, TF::TPUExecuteAndUpdateVariablesOp>(def_op))) {
      return true;
    }
    if (BlockArgument block_arg = mlir::dyn_cast<BlockArgument>(arg)) {
      // Skip the folding logic if the block argument is not from the function
      // arguments. This can happen when the argument is from a while loop.

                           result.getResultNumber());
    return std::string(result.getOwner()->getName().getStringRef());
  }
  // Use the ASM syntax for BlockArgument
  if (auto arg = mlir::dyn_cast<mlir::BlockArgument>(val)) {
    return "arg" + std::to_string(arg.getArgNumber());
  }
  return "";
}

std::string OpOrArgStripNameMapper::GetName(OpOrVal op_or_val) {

// replicate op is being modified.
llvm::ArrayRef<BlockArgument> ReplicateOp::GetReplicatedBlockArguments() {
  return GetBody().getArguments().drop_back(GetNumPackedBlockArguments());
}

// Returns the packed block arguments. A copy should be made if the replicate op
// is being modified.
llvm::ArrayRef<BlockArgument> ReplicateOp::GetPackedBlockArguments() {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/util/device_name_utils.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

namespace {

using mlir::BlockArgument;
using mlir::failure;
using mlir::LogicalResult;
using mlir::Operation;
using mlir::OperationPass;
using mlir::OpOperand;
using mlir::StringAttr;
using mlir::success;
using mlir::Value;

                      dst_index);
    }
    return absl::OkStatus();
  }

  auto input_arg = mlir::cast<BlockArgument>(src);
  auto input_node_it = args_.find(input_arg);
  TF_RET_CHECK(input_node_it != args_.end())
      << "Use of BlockArgument encounted before def!";
  // For argument, there is only one result output, so the index is always 0.