Type GetEmbeddingLookupShape(Value lookup, Value value) {
  auto lookup_type = mlir::cast<ShapedType>(lookup.getType());
  if (!lookup_type.hasStaticShape()) {
    return nullptr;
  }
  auto value_type = mlir::cast<ShapedType>(value.getType());
  if (!value_type.hasStaticShape() || value_type.getRank() != 2) {
    return nullptr;
  }
  SmallVector<int64_t> new_shape = {lookup_type.getNumElements(),