for (BlockArgument& func_arg : func_args) {
    auto resource_type = mlir::dyn_cast<TF::ResourceType>(
        getElementTypeOrSelf(func_arg.getType()));
    if (!resource_type) continue;
    if (failed(ValidateResourceArgument(function, func_arg, resource_type)))
      return failure();

    add_resource_argument(func_arg, resource_type);
  }

static Type GetResourceSubtypeOrDefault(Value resource, Type element_type) {
  auto resource_type = resource.getType()
                           .cast<TensorType>()
                           .getElementType()
                           .cast<ResourceType>();
  if (resource_type.getSubtypes().size() == 1)
    return resource_type.getSubtypes().front();

  return UnrankedTensorType::get(element_type);
}

    auto it =
        shared_name_to_arg_idx.insert({shared_name, block.getNumArguments()});
    if (it.second) {
      auto resource_type = op.getResult(0).getType();
      op.getResult(0).replaceAllUsesWith(
          block.addArgument(resource_type, op.getLoc()));
      AddEntryFunctionInput(
          absl::StrCat("hash_table_", it.first->getValue(), ":0"), target_func);

    const auto resource_type = RankedTensorType::get(
        /*shape=*/{}, /*elementType=*/TF::ResourceType::get(
            /*subtypes=*/llvm::ArrayRef<TensorType>{const_op.getType()},
            builder.getContext()));
    auto var_handle_op =
        builder.create<TF::VarHandleOp>(const_op.getLoc(),
                                        /*resource=*/resource_type,

      auto resource_type =
          mlir::TF::ResourceType::get({resource_subtype}, builder.getContext());

      auto tensor_type = mlir::cast<mlir::TensorType>(mlir_arg.getType());
      if (tensor_type.hasRank()) {
        mlir_arg.setType(
            GetTypeFromTFTensorShape(tensor_type.getShape(), resource_type));
      } else {

  node_def->set_op(FunctionLibraryDefinition::kArgOp);

  mlir::TensorType arg_type = mlir::cast<mlir::TensorType>(arg.getType());
  if (auto resource_type =
          mlir::dyn_cast<mlir::TF::ResourceType>(arg_type.getElementType())) {
    llvm::ArrayRef<mlir::TensorType> subtypes = resource_type.getSubtypes();
    if (!subtypes.empty()) {
      AttrValue handle_dtypes_attr;
      AttrValue handle_shapes_attr;

  node_def->set_op(FunctionLibraryDefinition::kArgOp);

  mlir::TensorType arg_type = mlir::cast<mlir::TensorType>(arg.getType());
  if (auto resource_type =
          mlir::dyn_cast<mlir::TF::ResourceType>(arg_type.getElementType())) {
    llvm::ArrayRef<mlir::TensorType> subtypes = resource_type.getSubtypes();
    if (!subtypes.empty()) {
      AttrValue handle_dtypes_attr;
      AttrValue handle_shapes_attr;

    return op.emitOpError("requires input to have one resource");

  auto resource_type = input_type.getElementType().cast<TF::ResourceType>();
  auto subtypes = resource_type.getSubtypes();
  switch (subtypes.size()) {
    case 1:
      return VerifyShapeOperandAndResult(
          op, resource_type.getSubtypes().front(), op.getType());
    case 0:

bool IsResource(Value value) {
  return mlir::isa<TF::ResourceType>(getElementTypeOrSelf(value.getType()));
}

// Get the type of the data contained in a resource. Returns null if there is
// no single type in the resource.
Type GetResourceSubtype(Value value) {
  auto resource_type =
      mlir::dyn_cast<TF::ResourceType>(getElementTypeOrSelf(value.getType()));
  auto subtypes = resource_type.getSubtypes();

      [{ mlir::TF::ShapeAttr::get($_ctxt, $_self) }]>;

  let extraClassDeclaration = [{
    TensorType resource_subtype() { return resource_type().getSubtypes()[0]; }

    ResourceType resource_type() {
      return getElementTypeOrSelf(getResource()).cast<TF::ResourceType>();
    }
  }];

  let hasVerifier = 1;
}

def TF_XlaShardingOp : TF_Op<"XlaSharding", [Pure, TF_NoConstantFold]> {