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]> {

namespace mlir {
namespace TFDevice {
namespace {

bool IsResourceType(Type val_type) {
  if (auto tensor_type = mlir::dyn_cast<mlir::TensorType>(val_type)) {
    if (mlir::isa<TF::ResourceType>(tensor_type.getElementType())) {
      return true;
    }
  }
  return false;
}

struct EmbeddingPipeliningPass
    : public ::impl::EmbeddingPipeliningPassBase<EmbeddingPipeliningPass> {