ops_disk doctrings

docs/_source/tinygrad-runtime-ops_disk.rst

@@ -0,0 +1,12 @@
+tinygrad runtime.ops_disk
+-------------------------
+
+.. note:: You likely want the upstream tinygrad, not tinygrab.
+          Tinygrab contains AI generated docstrings for a tinygrad snapshot.
+          Upstream: https://tinygrad.org
+
+.. automodule:: tinygrad.runtime.ops_disk
+   :members:
+   :undoc-members:
+   :show-inheritance:
+

docs/_source/tinygrad.rst

@@ -22,6 +22,7 @@ tinygrad
    tinygrad-runtime-ops_clang
    tinygrad-runtime-ops_cpu
    tinygrad-runtime-ops_cuda
+   tinygrad-runtime-ops_disk
    :maxdepth: 3
    :caption: Contents:
 

tinygrad/runtime/ops_disk.py

@@ -12,27 +12,77 @@ from tinygrad.shape.view import strides_for_shape
 
 
 class UnderlyingDiskBuffer:
+    """
+    This class represents an underlying disk buffer. It is initialized with a file descriptor (fd) and memory (mem).
+
+    Attributes:
+        fd (file descriptor): The file descriptor of the disk buffer.
+        mem (memory): The in-memory data associated with the disk buffer.
+    """
     def __init__(self, fd, mem):
+        """
+        Constructs all the necessary attributes for the underlying disk buffer.
+
+        Parameters:
+            fd (file descriptor): The file descriptor of the disk buffer.
+            mem (memory): The in-memory data associated with the disk buffer.
+        """
         self.fd, self.mem = fd, mem
 
     def __del__(self):
+        """
+        Closes the file descriptor if it exists when the object is destroyed.
+        """
         if self.fd:
             self.fd.close()
 
 
 class DiskBuffer:
+    """
+    Class for handling disk buffer operations.
+
+    Attributes:
+        ud (UnderlyingDiskBuffer): Underlying disk buffer.
+        size (int): Size of the buffer.
+        dtype (DType): Data type of the buffer elements. Default is dtypes.uint8.
+        offset (int): Offset in the buffer. Default is 0.
+    """
     def __init__(
         self, ud: UnderlyingDiskBuffer, size: int, dtype: DType = dtypes.uint8, offset=0
     ):
         self.ud, self.size, self.dtype, self.offset = ud, size, dtype, offset
 
     def __repr__(self):
+        """
+        Return a string representation of the DiskBuffer object.
+
+        Returns:
+            str: String representation of the DiskBuffer object containing its size, data type and offset.
+        """
         return f"<DiskBuffer size={self.size} dtype={self.dtype} offset={self.offset}>"
 
     def cast(self, arg: Tuple[DType, bool]):
+        """
+        Cast the DiskBuffer to a new data type.
+
+        Args:
+            arg (Tuple[DType, bool]): A tuple containing the new data type and a boolean value.
+
+        Returns:
+            DiskBuffer: A new DiskBuffer with the casted data type.
+        """
         return DiskBuffer(self.ud, self.size, arg[0], offset=self.offset)
 
     def as_strided(self, arg):
+        """
+        Create a view of the original buffer with new dimensions and strides.
+
+        Args:
+            arg: A tuple containing the new shape, calculated strides and an offset.
+
+        Returns:
+            DiskBuffer: A new DiskBuffer with the reshaped data.
+        """
         assert strides_for_shape(arg[0]) == arg[1], "disk tensors don't support strides"
         return DiskBuffer(
             self.ud,
@@ -42,6 +92,12 @@ class DiskBuffer:
         )
 
     def _buf(self) -> memoryview:
+        """
+        Return a memory view of the buffer.
+
+        Returns:
+            memoryview: A memory view of the underlying buffer.
+        """
         return memoryview(self.ud.mem)[
             self.offset : self.offset + self.size * self.dtype.itemsize
         ]
@@ -56,10 +112,31 @@ MAP_LOCKED, MAP_POPULATE = 0x2000, 0x008000
 
 
 class DiskAllocator(Allocator):
+    """
+    DiskAllocator class for allocating disk space.
+
+    Attributes:
+        device (str): The device to be used for allocation.
+    """
     def __init__(self, device):
+        """
+        Initializes the DiskAllocator object.
+
+        Parameters:
+            device (str): The device to use for disk allocation.
+        """
         self.device = device
 
     def _alloc(self, size):
+        """
+        Allocates memory from the device.
+
+        Parameters:
+            size (int): Size of the memory to allocate in bytes.
+
+        Returns:
+            DiskBuffer: A DiskBuffer object representing the allocated memory.
+        """
         if str(self.device).startswith("shm:"):
             if OSX:
                 with open(f"/tmp/shm_{self.device[4:]}", "w+b") as f:
@@ -82,12 +159,35 @@ class DiskAllocator(Allocator):
         return DiskBuffer(buf, size)
 
     def as_buffer(self, src: DiskBuffer):
+        """
+        Returns the underlying buffer of a DiskBuffer object.
+
+        Parameters:
+            src (DiskBuffer): The DiskBuffer object to get the buffer from.
+
+        Returns:
+            UnderlyingDiskBuffer: The underlying buffer of the DiskBuffer object.
+        """
         return src._buf()
 
     def copyin(self, dest: DiskBuffer, src: memoryview):
+        """
+        Copies data from a memoryview object to a DiskBuffer object.
+
+        Parameters:
+            dest (DiskBuffer): The destination DiskBuffer object.
+            src (memoryview): The source memoryview object.
+        """
         dest._buf()[:] = src
 
     def copyout(self, dest: memoryview, src: DiskBuffer):
+        """
+        Copies data from a DiskBuffer object to a memoryview object.
+
+        Parameters:
+            dest (memoryview): The destination memoryview object.
+            src (DiskBuffer): The source DiskBuffer object.
+        """
         if src.ud.fd is not None:
             src.ud.fd.seek(src.offset)
             src.ud.fd.readinto(dest)
@@ -96,5 +196,16 @@ class DiskAllocator(Allocator):
 
 
 class DiskDevice(Interpreted):
+    """
+    Initialize a new DiskDevice object.
+
+    Parameters:
+        device (str): The device identifier for the disk to be allocated.
+
+    Attributes:
+        super (Interpreted): Inherit from the Interpreted class.
+        DiskAllocator (DiskAllocator): Initialize a new DiskAllocator object with the given device.
+        disk_fxn_for_op (function): The function to be used for performing operations on the disk.
+    """
     def __init__(self, device):
         super().__init__(DiskAllocator(device[5:]), disk_fxn_for_op)