Method fs_load

tinygrad/tensor.py:414–440

Load a tensor from storage. self should be a tensor of the hash to load

(self, size:int)

  CHUNK_SIZE = 2**20
  def fs_load(self, size:int) -> Tensor:
    """
    Load a tensor from storage.

    self should be a tensor of the hash to load
    """
    # TODO: this should work locally as well
    assert self.dtype == dtypes.uint8, "hash is expected to be uint8"
    h = self.contiguous().flatten()
    assert h.shape[0] == 16, "expected hash"

    base_chunks = math.ceil(size / Tensor.CHUNK_SIZE)
    tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
    data, level_chunks = h, 0
    for i in reversed(range(tree_depth + 1)):
      data = data.to("tinyfs:load")

      # if not last level, its still hashes
      if i > 0 or tree_depth == 0:
        level_chunks = max(1, math.ceil(base_chunks / (Tensor.CHUNK_SIZE // 16)**(i-1)))
        pad_amt = 16 * level_chunks
      else: pad_amt = Tensor.CHUNK_SIZE * level_chunks
      if (tsize := data.shape[0]) < pad_amt: data = data.pad((0, pad_amt - tsize))
      data = data[:pad_amt].contiguous()
      if i != 0: data = data.to(self.device)

    return data[:size]

  def fs_store(self) -> Tensor:
    """