}

// Checks the result Variant type to infer the element shape if fully defined.
// If the Variant type has multiple subtypes or does not have static shape,
// return error.
LogicalResult GetElementShapeFromResultType(
    Type type, llvm::SmallVector<int64_t, 8>* shape) {
  auto variant_type =
      mlir::dyn_cast<TF::VariantType>(getElementTypeOrSelf(type));