lox.load#

lox.load(path, key=None, result_shape_dtypes=None, argnames=None)#

Load data from a specified path. Each file in the directory is loaded into a dictionary with the filename (without extension) as the key.

Parameters:
  • path (StringArray) – The path from which the data will be loaded.

  • result_shape_dtypes (Any, optional) – The expected shape and dtype of the loaded data.

  • argnames (Iterable[str], optional) – Specific argument names to load. If None, all files in the directory are loaded.

  • key (jax.Array, optional) – An optional key to differentiate data when loading.

Returns:

The loaded data.

Return type:

dict[str, Any]