Store a tensor to storage.
(self)
| 440 | return data[:size] |
| 441 | |
| 442 | def fs_store(self) -> Tensor: |
| 443 | """ |
| 444 | Store a tensor to storage. |
| 445 | """ |
| 446 | # TODO: this should work locally as well |
| 447 | data = self.contiguous().flatten().bitcast(dtypes.uint8) |
| 448 | |
| 449 | # pad to a multiple of 1mb |
| 450 | if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE)) |
| 451 | size = data.shape[0] |
| 452 | |
| 453 | base_chunks = math.ceil(size / Tensor.CHUNK_SIZE) |
| 454 | tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16)) |
| 455 | |
| 456 | to_device = "CPU" if isinstance(self.device, str) and self.device.startswith("DISK") else self.device |
| 457 | |
| 458 | level_chunks = base_chunks |
| 459 | for _ in range(tree_depth + 1): |
| 460 | data = data.to("tinyfs:store")[:level_chunks * 16].contiguous().to(to_device) |
| 461 | if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE)) |
| 462 | level_chunks = math.ceil(data.shape[0] / Tensor.CHUNK_SIZE) |
| 463 | |
| 464 | return data[:16].contiguous() |
| 465 | |
| 466 | @staticmethod |
| 467 | def from_uop(y:UOp, **kwargs) -> Tensor: |