def wrap_with_cached_feature(
    cached_feature_type,
    fallback_features: Union[Feature, Dict[FeatureKey, Feature]],
    max_cache_size_in_bytes: int,
    *args,
    **kwargs,
) -> Union[Feature, Dict[FeatureKey, Feature]]:
    """Wraps the given features with the given cached feature type using
    a single cache instance."""
    if not isinstance(fallback_features, dict):
        assert isinstance(fallback_features, Feature)
        return wrap_with_cached_feature(
            cached_feature_type,
            {"a": fallback_features},
            max_cache_size_in_bytes,
            *args,
            **kwargs,
        )["a"]
    row_bytes = None
    cache = None
    wrapped_features = {}
    offset = 0
    for feature_key, fallback_feature in fallback_features.items():
        # Fetching the feature dimension from the underlying feature.
        feat0 = fallback_feature.read(torch.tensor([0]))
        if row_bytes is None:
            row_bytes = feat0.nbytes
        else:
            assert (
                row_bytes == feat0.nbytes
            ), "The # bytes of a single row of the features should match."
        cache_size = bytes_to_number_of_items(max_cache_size_in_bytes, feat0)
        if cache is None:
            cache = cached_feature_type._cache_type(
                cache_shape=(cache_size,) + feat0.shape[1:],
                dtype=feat0.dtype,
                *args,
                **kwargs,
            )
        wrapped_features[feature_key] = cached_feature_type(
            fallback_feature, cache=cache, offset=offset
        )
        offset += fallback_feature.count()

    return wrapped_features
Wraps the given features with the given cached feature type using
    a single cache instance.
python/dgl/graphbolt/feature_store.py