feat!: remove NumPy and use our Rust backend

.github/workflows/main.yaml

@@ -55,7 +55,9 @@ jobs:
         run: uv sync --locked --dev
 
       - name: Ruff
-        run: uv run ruff check .
+        run: |
+          uv run ruff check .
+          uv run ruff format --check .
 
       - name: Pyright
         run: uv run pyright

Cargo.toml

@@ -10,9 +10,12 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 pyo3 = { version = "0.24.1", features = ["extension-module"] }
 rand = { version = "0.8.5" }
+rand_distr = "0.4"
 thiserror = "1"
 
 [features]
+default = ["python"]
+python = []
 abi3 = ["pyo3/abi3-py37", "generate-import-lib"]
 generate-import-lib = ["pyo3/generate-import-lib"]
 

cranberry/cranberry.pyi

@@ -6,17 +6,30 @@ class StorageView:
   def full(value: float, size: int, device: str) -> StorageView: ...
   @staticmethod
   def from_vec(vec: Union[list[float], np.ndarray], device: str) -> StorageView: ...
+  @staticmethod
+  def zeros(shape: list[int], device: str = ...) -> StorageView: ...
+  @staticmethod
+  def ones(shape: list[int], device: str = ...) -> StorageView: ...
+  @staticmethod
+  def randn(shape: list[int], device: str = ..., seed: int | None = ...) -> StorageView: ...
+  @staticmethod
+  def uniform(shape: list[int], low: float, high: float, device: str = ..., seed: int | None = ...) -> StorageView: ...
   def len(self) -> int: ...
+  def shape(self) -> list[int]: ...
   def to_vec(self) -> list[float]: ...
   def slice(self, offset: int, size: int) -> StorageView: ...
   def reshape(self, shape: list[int]) -> StorageView: ...
   def expand(self, shape: list[int]) -> StorageView: ...
   def permute(self, dims: list[int]) -> StorageView: ...
+  def contiguous(self) -> StorageView: ...
   def neg(self) -> StorageView: ...
   def sqrt(self) -> StorageView: ...
+  def relu(self) -> StorageView: ...
   def exp(self) -> StorageView: ...
   def log(self) -> StorageView: ...
   def add(self, other: StorageView) -> StorageView: ...
   def sub(self, other: StorageView) -> StorageView: ...
   def mul(self, other: StorageView) -> StorageView: ...
   def div(self, other: StorageView) -> StorageView: ...
+  def sum(self, dim: int | None, keepdim: bool = ...) -> StorageView: ...
+  def max(self, dim: int | None, keepdim: bool = ...) -> StorageView: ...

cranberry/features/datasets.py

@@ -8,20 +8,39 @@
 import urllib.request
 import sys
 
-import numpy as np
+from typing import TYPE_CHECKING, Any
+
 from cranberry import Tensor
 
+if TYPE_CHECKING:  # pragma: no cover - typing only
+  pass
+
+try:  # pragma: no cover - optional dependency
+  import numpy as np  # type: ignore[assignment]
+except ImportError:  # pragma: no cover
+  np = None  # type: ignore[assignment]
+
+
+def _require_numpy() -> Any:
+  if np is None:  # pragma: no cover - optional path
+    raise RuntimeError("NumPy is required for dataset utilities. Install 'cranberry[numpy]' to enable them.")
+  return np
+
+
 # Platform detection without psutil
 OSX = sys.platform == "darwin"
 
 # Optional tqdm progress bar
 try:
   from tqdm import tqdm as _tqdm  # type: ignore
+
   tqdm = _tqdm  # noqa: N802 - keep name compatibility
 except Exception:  # pragma: no cover - fallback path
+
   class _NoopTqdm:
     def __init__(self, *a, **kw):
       pass
+
     def update(self, n: int):
       pass
 
@@ -60,10 +79,11 @@ def fetch(
 
 
 def _fetch_mnist(file, offset):
+  np_mod = _require_numpy()
   return Tensor(
-    np.frombuffer(
+    np_mod.frombuffer(
       gzip.open(fetch("https://storage.googleapis.com/cvdf-datasets/mnist/" + file)).read()[offset:],
-      dtype=np.uint8,
+      dtype=np_mod.uint8,
     )
   )
 

cranberry/nn/__init__.py

@@ -7,16 +7,20 @@
 
 class Module:
   @abstractmethod
-  def __call__(self, x: Tensor) -> Tensor: pass
+  def __call__(self, x: Tensor) -> Tensor:
+    pass
 
   @abstractmethod
-  def parameters(self) -> List[Tensor]: pass
+  def parameters(self) -> List[Tensor]:
+    pass
 
 
 class ReLU(Module):
-  def __call__(self, x: Tensor) -> Tensor: return x.relu()
+  def __call__(self, x: Tensor) -> Tensor:
+    return x.relu()
 
-  def parameters(self) -> List[Tensor]: return []
+  def parameters(self) -> List[Tensor]:
+    return []
 
 
 class Linear(Module):
@@ -38,7 +42,8 @@ def __init__(self, *layers: Module):
     self.layers = layers
 
   def __call__(self, x: Tensor) -> Tensor:
-    for layer in self.layers: x = layer(x)
+    for layer in self.layers:
+      x = layer(x)
     return x
 
   def parameters(self) -> List[Tensor]:

cranberry/ops.py

@@ -3,19 +3,41 @@
 
 
 class UnaryOps(Enum):
-  NEG = auto(); SQRT = auto(); RELU = auto(); EXP = auto(); LOG = auto() # noqa: E702
-  def __repr__(self): return f"{self.name.lower()}"
+  NEG = auto()
+  SQRT = auto()
+  RELU = auto()
+  EXP = auto()
+  LOG = auto()
+
+  def __repr__(self):
+    return f"{self.name.lower()}"
+
 
 class BinaryOps(Enum):
-  ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto() # noqa: E702
-  def __repr__(self): return f"{self.name.lower()}"
+  ADD = auto()
+  SUB = auto()
+  MUL = auto()
+  DIV = auto()
+
+  def __repr__(self):
+    return f"{self.name.lower()}"
+
 
 class ReduceOps(Enum):
-  SUM = auto(); MAX = auto() # noqa: E702
-  def __repr__(self): return f"{self.name.lower()}"
+  SUM = auto()
+  MAX = auto()
+
+  def __repr__(self):
+    return f"{self.name.lower()}"
+
 
 class MovementOps(Enum):
-  RESHAPE = auto(); EXPAND = auto(); PERMUTE = auto() # noqa: E702
-  def __repr__(self): return f"{self.name.lower()}"
+  RESHAPE = auto()
+  EXPAND = auto()
+  PERMUTE = auto()
+
+  def __repr__(self):
+    return f"{self.name.lower()}"
+
 
 Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, None]

cranberry/optim/__init__.py

@@ -1,4 +1,4 @@
-from cranberry import Tensor
+from cranberry import StorageView, Tensor
 from typing import List
 
 
@@ -11,16 +11,29 @@ def __init__(self, params: List[Tensor], lr: float):
     self._params, self._lr = params, lr
 
   def zero_grad(self):
-    for p in self._params: p._grad.fill(0.0)
+    for p in self._params:
+      p.zero_grad()
 
   def step(self):
-    for p in self._params: p._data -= self._lr * p._grad
+    for p in self._params:
+      grad_storage = p.grad_storage()
+      if grad_storage is None:
+        continue
+      grad_view = grad_storage.contiguous()
+      shape_list = list(p.shape)
+      scale = StorageView.full(float(self._lr), max(p.num_elements(), 1), p.device)
+      scale = scale.reshape(shape_list) if shape_list else scale.reshape([])
+      scaled_grad = grad_view.mul(scale)
+      new_data = p.data_storage().contiguous().sub(scaled_grad)
+      p.set_data_storage(new_data)
 
   @property
-  def lr(self): return self._lr
+  def lr(self):
+    return self._lr
 
   @lr.setter
-  def lr(self, lr: float): self._lr = lr
+  def lr(self, lr: float):
+    self._lr = lr
 
 
 # https://pytorch.org/docs/stable/generated/torch.optim.Adam.html