How Tensor Bytes Move

How Tensor Bytes Move A data-flow diagram generated by Archify. 01 / Outside 02 / Bridge in 03 / Storage 04 / Exit Host values · list · scalar · DLPack · 01 / Outside Host values list · scalar · DLPack HF Hub · lucid-dl · 01 / Outside HF Hub lucid-dl Dataset · __getitem__ · 01 / Outside Dataset __getitem__ lucid.tensor · converters.py · 02 / Bridge in lucid.tensor converters.py weights._hub · cache · SHA-256 · 02 / Bridge in weights._hub cache · SHA-256 DataLoader · default_collate · 02 / Bridge in DataLoader default_collate GpuStorage · mlx::core::array · 03 / Storage GpuStorage mlx::core::array CpuStorage · 64B aligned · pool · 03 / Storage CpuStorage 64B aligned · pool SharedStorage · Metal shared buffer · 03 / Storage SharedStorage Metal shared buffer Graph export · compile · coreml · 04 / Exit Graph export compile · coreml lucid.save · pickle · tensor_v3 · 04 / Exit lucid.save pickle · tensor_v3 NumPy/DLPack · numpy() · __dlpack__ · 04 / Exit NumPy/DLPack numpy() · __dlpack__ list · ndarray always copied safetensors SHA-256 checked samples Python objects from_bytes load_safetensors stack to('metal') cpu_view MTLBuffer zero-copy to_bytes copied numpy() zero-copy view Legend primary data zero-copy data store data flow

Copying is the default

  • • tensor, from_numpy and as_tensor all copy, despite from_numpy's 'shares memory' docstring
  • • Eager reshape, permute and slice allocate a new buffer on CPU

Where zero-copy happens

  • • A kDLMetal DLPack capsule is adopted as an MLX array
  • • CPU numpy(), SharedStorage cpu_view, compile feeds and Core ML inputs

GPU is lazy

  • • GpuStorage stays an MLX graph until evaluated
  • • item, numpy, repr and save force eval plus a CPU download