Function dump_tensor_cache

python/tvm/contrib/tvmjs.py:202–330

Dump parameters to Tensor cache. Parameters ---------- params: Union[ Mapping[str, Union[np.ndarray, tvm.runtime.Tensor]], Iterator[Tuple[str, Union[np.ndarray, tvm.runtime.Tensor]]], ] The parameter dictionary or generator cache_dir: str The pat

(
    params: Mapping[str, np.ndarray | tvm.runtime.Tensor]
    | Iterator[tuple[str, np.ndarray | tvm.runtime.Tensor]],
    cache_dir: str,
    encode_format="f32-to-bf16",
    meta_data=None,
    shard_cap_mb=32,
    show_progress: bool = True,
    update_if_exists: bool = False,
)

def dump_tensor_cache(
    params: Mapping[str, np.ndarray | tvm.runtime.Tensor]
    | Iterator[tuple[str, np.ndarray | tvm.runtime.Tensor]],
    cache_dir: str,
    encode_format="f32-to-bf16",
    meta_data=None,
    shard_cap_mb=32,
    show_progress: bool = True,
    update_if_exists: bool = False,
):
    """Dump parameters to Tensor cache.

    Parameters
    ----------
    params: Union[
        Mapping[str, Union[np.ndarray, tvm.runtime.Tensor]],
        Iterator[Tuple[str, Union[np.ndarray, tvm.runtime.Tensor]]],
    ]
        The parameter dictionary or generator

    cache_dir: str
        The path to the cache

    encode_format: {"f32-to-bf16", "raw"}
        Encoding format.

    meta_data: json-compatible-struct or Callable[[], Any]
        Extra meta_data to be stored in the cache json file,
        or a callable that returns the metadata.

    shard_cap_mb: int
        Maxinum number of MB to be kept per shard

    show_progress: bool
        A boolean indicating if to show the dump progress.

    update_if_exists: bool
        If the cache already exists, update the cache. When set to False, it will overwrite the
        existing files.
    """
    if encode_format not in ("raw", "f32-to-bf16"):
        raise ValueError(f"Invalie encode_format {encode_format}")

    records = []
    from_generator = isinstance(params, GeneratorType)
    total_bytes = 0
    counter = 0
    max_out_length = 0

    if not os.path.exists(cache_dir):
        os.makedirs(cache_dir)

    f32_to_bf16_triggered = False

    print(f"Start storing to cache {cache_dir}")
    shard_cap_nbytes = shard_cap_mb * (1 << 20)

    nd_cache_json = os.path.join(cache_dir, "tensor-cache.json")