return UnrankedTensorType::get(
        variant_ty.getSubtypes()[0].getElementType());
  }
  Type value_type = value.getType();
  Type element_type;
  variant_ty = value_type.dyn_cast<TF::VariantType>();
  if (variant_ty && !variant_ty.getSubtypes().empty()) {
    element_type = variant_ty.getSubtypes()[0].getElementType();
  } else {
    element_type = getElementTypeOrSelf(value_type);
  }

      if (mlir::isa<QuantizedType>(type.getElementType())) {
        has_quantized_types = true;
      }
    }
  }
  for (Value output : call_op.getOutput()) {
    if (auto type = mlir::dyn_cast<TensorType>(output.getType())) {
      if (mlir::isa<QuantizedType>(type.getElementType())) {
        has_quantized_types = true;
      }
    }
  }

bool IsI8ToF32Cast(stablehlo::ConvertOp convert_op) {
  const bool is_i8_operand =
      convert_op.getOperand().getType().getElementType().isInteger(/*width=*/8);
  const bool is_f32_result =
      mlir::isa<Float32Type>(convert_op.getResult().getType().getElementType());
  return is_i8_operand && is_f32_result;
}

// Tests whether a `stablehlo::ConvertOp` is a i8 -> i32 cast.

             .getElementType()
             .isF32()) {
      return failure();
    }
    MLIRContext *context = rewriter.getContext();
    llvm::SmallVector<Value, 2> operands{op.getA(), op.getB()};
    for (Value &operand : operands) {
      TensorType tensor_type = mlir::cast<TensorType>(operand.getType());
      Type element_type = tensor_type.getElementType();

    // TODO(antiagainst): also handle the case of tf.Add(tf.[op], <bias>)

    // Get a splat zero tensor with the expected dimension for the bias tensor
    auto elem_type = filter_type.getElementType();
    auto bias_dim = static_cast<const ConcreteType *>(this)->getBiasDim(
        filter_type.getShape());
    auto bias_type =
        tensorflow::GetTypeFromTFTensorShape({bias_dim}, elem_type);

  CPred<"$_self.isa<ElementsAttr>() && $_self.cast<ElementsAttr>().getShapedType().getElementType().isF32()">,
        "32 bit float constant tensor">;

// Checks if the param passed is a float ElementsAttr.
def FloatElementsAttr : ElementsAttrBase<
  CPred<"$_self.isa<ElementsAttr>() && $_self.cast<ElementsAttr>().getShapedType().getElementType().isa<FloatType>()">,
        "float constant tensor">;

    bool truncate = shaped_type.getElementType().getIntOrFloatBitWidth() == 64;
    TF_ASSIGN_OR_RETURN(value,
                        tfl::ConvertIntBuffer(shaped_type, buffer, truncate));
    TF_ASSIGN_OR_RETURN(
        mlir::quant::QuantizedType type,
        tfl::GetQuantizedType(tensor, builder, /*is_constant=*/true,
                              /*storage_type=*/value.getElementType()));
    shaped_type = shaped_type.clone(type);

        given:
        addToContainer(b)
        container.all {
            seen << it
        }

        when:
        container.addAllLater(provider)

        then:
        _ * provider.getElementType() >> getType()
        _ * provider.size() >> 2
        _ * provider.calculateValue(_) >> ValueSupplier.Value.of([a, d])
        0 * _
        seen == [b, a, d]
    }

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