forge-dl 0.1.0__py3-none-any.whl

forge/__init__.py

@@ -0,0 +1,19 @@
+"""
+Forge: a machine learning library and deep learning compiler written from
+scratch in Python.
+
+The core has no dependencies. Every operation, from the autograd engine to the
+attention mechanism, is implemented from first principles.
+"""
+
+from forge.tensor import Tensor
+from forge.dtype import float32, float64, int32, int64
+from forge.serialization import save_model, load_model
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Tensor",
+    "float32", "float64", "int32", "int64",
+    "save_model", "load_model",
+]

forge/autograd/__init__.py

@@ -0,0 +1,20 @@
+"""
+The autograd engine: reverse-mode automatic differentiation on a define-by-run
+graph.
+
+Every operation records itself as it executes, and backward() walks that record
+in reverse, applying the chain rule at each node.
+"""
+
+from forge.autograd.engine import Function
+from forge.autograd.operations import (
+    Add, Mul, Sub, Pow, Neg, Sum, Mean, Matmul,
+    ReLU, Sigmoid, Tanh, Log, Clamp, Softmax, Exp,
+)
+from forge.autograd.grad_check import grad_check
+
+__all__ = [
+    "Function", "grad_check",
+    "Add", "Mul", "Sub", "Pow", "Neg", "Sum", "Mean", "Matmul",
+    "ReLU", "Sigmoid", "Tanh", "Log", "Clamp", "Softmax", "Exp",
+]

forge/autograd/accelerate_backend.py

@@ -0,0 +1,70 @@
+import array
+import ctypes
+
+try:
+    _lib = ctypes.CDLL("/System/Library/Frameworks/Accelerate.framework/Accelerate")
+    
+    cblas_sgemm = _lib.cblas_sgemm
+
+    float_pointer = ctypes.POINTER(ctypes.c_float)
+
+    cblas_sgemm.argtypes = [
+        ctypes.c_int,   # storage order (row-major vs column-major)
+        ctypes.c_int,   # whether to transpose the left matrix
+        ctypes.c_int,   # whether to transpose the right matrix
+        ctypes.c_int,   # number of rows in the result
+        ctypes.c_int,   # number of columns in the result
+        ctypes.c_int,   # the shared inner dimension being summed over
+        ctypes.c_float, # alpha: the product is scaled by this
+        float_pointer,  # pointer to the left matrix's data
+        ctypes.c_int,   # row stride of the left matrix
+        float_pointer,  # pointer to the right matrix's data
+        ctypes.c_int,   # row stride of the right matrix
+        ctypes.c_float, # beta: the existing result is scaled by this
+        float_pointer,  # pointer to the output matrix's data
+        ctypes.c_int,   # row stride of the output matrix 
+    ]
+
+    cblas_sgemm.restype = None
+
+    ACCELERATE_AVAILABLE = True
+
+# Only occurs on non-macOS systems
+except Exception:
+    ACCELERATE_AVAILABLE = False
+
+CBLAS_ROW_MAJOR = 101
+CBLAS_NO_TRANSPOSE = 111
+
+BYTES_PER_FLOAT = 4
+
+def accelerate_matmul(left_data, right_data, left_rows, shared_dim, right_cols):
+    """
+    Multiplies two matrices on the CPU using Accelerate.
+
+    Returns a flat float32 array of the result.
+    """
+    result_data = array.array('f', bytes(BYTES_PER_FLOAT * left_rows * right_cols))
+
+    left_pointer = ctypes.cast(left_data.buffer_info()[0], float_pointer)
+    right_pointer = ctypes.cast(right_data.buffer_info()[0], float_pointer)
+    result_pointer = ctypes.cast(result_data.buffer_info()[0], float_pointer)
+
+    cblas_sgemm(
+        CBLAS_ROW_MAJOR,      # data is laid out row by row
+        CBLAS_NO_TRANSPOSE,   # do not transpose the left matrix
+        CBLAS_NO_TRANSPOSE,   # do not transpose the right matrix
+        left_rows,            # rows of the result
+        right_cols,           # columns of the result
+        shared_dim,           # the inner dimension summed over
+        1.0,                  # alpha: scale the product by 1.0
+        left_pointer,         # the left matrix
+        shared_dim,           # left's row stride
+        right_pointer,        # the right matrix
+        right_cols,           # right's row stride
+        0.0,                  # beta: ignore the zeroed existing result
+        result_pointer,       # where to write the answer
+        right_cols,           # result's row stride
+    )
+
+    return result_data

forge/autograd/engine.py

@@ -0,0 +1,14 @@
+class Function:
+    """Base class for all differentiable operations"""
+    def __init__(self):
+        self.inputs = []
+        self.saved_tensors = []
+
+    def save_for_backward(self, *tensors):
+        self.saved_tensors = list(tensors)
+
+    def forward(self, *args):
+        raise NotImplementedError
+
+    def backward(self, grad_output):
+        raise NotImplementedError

forge/autograd/fused_operations.py

@@ -0,0 +1,120 @@
+import array as _array
+import math
+from forge.autograd.engine import Function
+
+
+class FusedLinearReLU(Function):
+    """Fused matmul + bias add + ReLU in a single pass"""
+
+    def forward(self, x, weight_t, bias):
+        self.inputs = [x, weight_t, bias]
+        self.save_for_backward(x, weight_t, bias)
+        from forge.tensor import Tensor
+
+        m = x.shape[0]
+        n = x.shape[1]
+        p = weight_t.shape[1]
+
+        new_data = _array.array(x.dtype.typecode, [])
+
+        for i in range(m):
+            for j in range(p):
+                # Matmul
+                total = 0.0
+                for k in range(n):
+                    total += x._data[i * n + k] * weight_t._data[k * p + j]
+
+                # Add bias
+                total += bias._data[j]
+
+                # ReLU
+                if total < 0:
+                    total = 0.0
+
+                new_data.append(total)
+
+        result = Tensor.__new__(Tensor)
+        result._data = new_data
+        result.shape = (m, p)
+        result.dtype = x.dtype
+        result.requires_grad = False
+        result.grad = None
+        result._grad_fn = None
+
+        # Save output for backward (need to know where ReLU killed values)
+        self._output = result
+        return result
+
+    def backward(self, grad_output):
+        x, weight_t, bias = self.saved_tensors
+        output = self._output
+        from forge.tensor import Tensor, _broadcast_shape, _broadcast_data
+        from forge.autograd.operations import _unbroadcast
+
+        m = x.shape[0]
+        n = x.shape[1]
+        p = weight_t.shape[1]
+
+        # Apply ReLU mask to grad_output
+        relu_grad_data = _array.array(
+            grad_output.dtype.typecode,
+            [g if o > 0 else 0.0 for g, o in zip(grad_output._data, output._data)]
+        )
+        relu_grad = Tensor.__new__(Tensor)
+        relu_grad._data = relu_grad_data
+        relu_grad.shape = grad_output.shape
+        relu_grad.dtype = grad_output.dtype
+        relu_grad.requires_grad = False
+        relu_grad.grad = None
+        relu_grad._grad_fn = None
+
+        # Grad for x: relu_grad @ weight_t.T
+        # weight_t is (n, p), weight_t.T is (p, n)
+        grad_x_data = _array.array(x.dtype.typecode, [0.0] * (m * n))
+        for i in range(m):
+            for j in range(n):
+                total = 0.0
+                for k in range(p):
+                    total += relu_grad._data[i * p + k] * weight_t._data[j * p + k]
+                grad_x_data[i * n + j] = total
+
+        grad_x = Tensor.__new__(Tensor)
+        grad_x._data = grad_x_data
+        grad_x.shape = x.shape
+        grad_x.dtype = x.dtype
+        grad_x.requires_grad = False
+        grad_x.grad = None
+        grad_x._grad_fn = None
+
+        # Grad for weight_t: x.T @ relu_grad
+        grad_wt_data = _array.array(x.dtype.typecode, [0.0] * (n * p))
+        for i in range(n):
+            for j in range(p):
+                total = 0.0
+                for k in range(m):
+                    total += x._data[k * n + i] * relu_grad._data[k * p + j]
+                grad_wt_data[i * p + j] = total
+
+        grad_wt = Tensor.__new__(Tensor)
+        grad_wt._data = grad_wt_data
+        grad_wt.shape = weight_t.shape
+        grad_wt.dtype = weight_t.dtype
+        grad_wt.requires_grad = False
+        grad_wt.grad = None
+        grad_wt._grad_fn = None
+
+        # Grad for bias: sum relu_grad along rows
+        grad_bias_data = _array.array(bias.dtype.typecode, [0.0] * p)
+        for i in range(m):
+            for j in range(p):
+                grad_bias_data[j] += relu_grad._data[i * p + j]
+
+        grad_bias = Tensor.__new__(Tensor)
+        grad_bias._data = grad_bias_data
+        grad_bias.shape = bias.shape
+        grad_bias.dtype = bias.dtype
+        grad_bias.requires_grad = False
+        grad_bias.grad = None
+        grad_bias._grad_fn = None
+
+        return grad_x, grad_wt, grad_bias

forge/autograd/fusion.py

@@ -0,0 +1,40 @@
+from forge.nn.layers import Linear, ReLULayer, FusedLinearReLULayer
+
+
+def optimize_model(layers):
+    """
+    Takes a list of layers and returns an optimized list
+    with fused operations where possible.
+    """
+    optimized = []
+    i = 0
+
+    while i < len(layers):
+        # Pattern: Linear followed by ReLU
+        if (i + 1 < len(layers) and isinstance(layers[i], Linear) and isinstance(layers[i + 1], ReLULayer)):
+            linear = layers[i]
+            in_f = linear.weight.shape[1]
+            out_f = linear.weight.shape[0]
+
+            fused = FusedLinearReLULayer(in_f, out_f)
+
+            # Copy weights from the original linear layer
+            fused.weight._data = linear.weight._data.__class__(
+                linear.weight.dtype.typecode, linear.weight._data
+            )
+
+            # Copy bias — need to flatten from (1, out_f) to (out_f,)
+            if linear.bias is not None:
+                import array as _array
+                fused.bias._data = _array.array(
+                    linear.bias.dtype.typecode, list(linear.bias._data)
+                )
+
+            optimized.append(fused)
+            i += 2  # Skip both the Linear and ReLU
+
+        else:
+            optimized.append(layers[i])
+            i += 1
+
+    return optimized

forge/autograd/grad_check.py

@@ -0,0 +1,58 @@
+"""
+It is important to draw wisdom from many different places.
+If you take it from only one place, it becomes rigid and stale.
+    - Uncle Iroh
+"""
+
+def grad_check(func, inputs, eps=1e-3, tol=1e-2):
+    """
+    Using first principles for differentiation: df/dx = lim(h->0) [f(x + h) - f(x - h)] / 2h
+    We use a small value of h to approximate the limit.
+    """
+    for inp in inputs:
+        inp.grad = None
+
+    output = func(*inputs)
+    output.backward()
+
+    analytical_grads = [inp.grad for inp in inputs]
+
+    # Check each input
+    for idx, inp in enumerate(inputs):
+        analytical_grad = analytical_grads[idx]
+
+        # Compute numerical gradient for each element
+        for i in range(len(inp._data)):
+            original = inp._data[i]
+
+            # f(x + h)
+            inp._data[i] = original + eps
+            # Clear grads and recompute
+            for inp2 in inputs:
+                inp2.grad = None
+                inp2._grad_fn = None
+            out_plus = func(*inputs)
+
+            # f(x - h)
+            inp._data[i] = original - eps
+            for inp2 in inputs:
+                inp2.grad = None
+                inp2._grad_fn = None
+            out_minus = func(*inputs)
+
+            # Restore
+            inp._data[i] = original
+
+            # Numerical gradient
+            numerical = (out_plus._data[0] - out_minus._data[0]) / (2 * eps)
+            analytical = analytical_grad._data[i]
+
+            diff = abs(numerical - analytical)
+            if diff > tol:
+                raise AssertionError(
+                    f"Gradient check failed for input {idx}, element {i}: "
+                    f"numerical={numerical:.6f}, analytical={analytical:.6f}, "
+                    f"diff={diff:.6f}"
+                )
+
+    return True

forge/autograd/mps_backend.py

@@ -0,0 +1,52 @@
+import array
+
+try:
+    import Metal
+    import MetalPerformanceShaders as MPS
+
+    _DEVICE = Metal.MTLCreateSystemDefaultDevice()
+    _QUEUE = _DEVICE.newCommandQueue() if _DEVICE is not None else None
+
+    MPS_AVAILABLE = _DEVICE is not None and _QUEUE is not None
+    if MPS_AVAILABLE:
+        _SHARED = Metal.MTLResourceStorageModeShared
+except Exception:
+    _DEVICE = None
+    _QUEUE = None
+    MPS_AVAILABLE = False
+
+_FLOAT_BYTES = 4
+GPU_MIN_WORK = 200000
+
+def _descriptor(rows, cols):
+    return MPS.MPSMatrixDescriptor.matrixDescriptorWithRows_columns_rowBytes_dataType_(
+        rows, cols, cols * _FLOAT_BYTES, MPS.MPSDataTypeFloat32
+    )
+
+def mps_matmul(a_data, b_data, m, k, n):
+    buf_a = _DEVICE.newBufferWithBytes_length_options_(
+        a_data.tobytes(), m * k * _FLOAT_BYTES, _SHARED
+    )
+    buf_b = _DEVICE.newBufferWithBytes_length_options_(
+        b_data.tobytes(), k * n * _FLOAT_BYTES, _SHARED
+    )
+    buf_c = _DEVICE.newBufferWithLength_options_(m * n * _FLOAT_BYTES, _SHARED)
+
+    mat_a = MPS.MPSMatrix.alloc().initWithBuffer_descriptor_(buf_a, _descriptor(m, k))
+    mat_b = MPS.MPSMatrix.alloc().initWithBuffer_descriptor_(buf_b, _descriptor(k, n))
+    mat_c = MPS.MPSMatrix.alloc().initWithBuffer_descriptor_(buf_c, _descriptor(m, n))
+
+    kernel = MPS.MPSMatrixMultiplication.alloc().initWithDevice_transposeLeft_transposeRight_resultRows_resultColumns_interiorColumns_alpha_beta_(
+        _DEVICE, False, False, m, n, k, 1.0, 0.0
+    )
+
+    cmd = _QUEUE.commandBuffer()
+    kernel.encodeToCommandBuffer_leftMatrix_rightMatrix_resultMatrix_(
+        cmd, mat_a, mat_b, mat_c
+    )
+    cmd.commit()
+    cmd.waitUntilCompleted()
+
+    out = array.array('f')
+    out.frombytes(buf_c.contents().as_buffer(m * n * _FLOAT_BYTES))
+    return out