feat!: rewrite whole Rust backend

.github/workflows/main.yaml

@@ -28,7 +28,7 @@ jobs:
 
       - uses: dtolnay/rust-toolchain@nightly
         with:
-          toolchain: nightly-2024-08-15
+          toolchain: nightly
           components: rustfmt, clippy
       - uses: Swatinem/rust-cache@v2
 
@@ -68,7 +68,7 @@ jobs:
 
       - uses: dtolnay/rust-toolchain@nightly
         with:
-          toolchain: nightly-2024-08-15
+          toolchain: nightly
       - uses: Swatinem/rust-cache@v2
 
       - name: Run tests

Cargo.toml

@@ -5,11 +5,12 @@ edition = "2021"
 
 [lib]
 name = "cranberry"
-crate-type = ["cdylib"]
+crate-type = ["cdylib", "rlib"]
 
 [dependencies]
 pyo3 = { version = "0.24.1", features = ["extension-module"] }
 rand = { version = "0.8.5" }
+thiserror = "1"
 
 [features]
 abi3 = ["pyo3/abi3-py37", "generate-import-lib"]

cranberry/__init__.py

@@ -1,2 +1,2 @@
 from cranberry.tensor import Tensor as Tensor
-from .cranberry import StoragePtr as StoragePtr
+from .cranberry import StorageView as StorageView

cranberry/cranberry.pyi

@@ -1,29 +1,22 @@
 from typing import Union
 import numpy as np
 
-class StoragePtr:
+class StorageView:
   @staticmethod
-  def full(value: float, size: int, device: str) -> StoragePtr: ...
+  def full(value: float, size: int, device: str) -> StorageView: ...
   @staticmethod
-  def from_vec(vec: Union[list[float], np.ndarray], device: str) -> StoragePtr: ...
-  @staticmethod
-  def neg(a: StoragePtr, b: StoragePtr, idx_a: int, idx_b: int, size: int): ...
-  @staticmethod
-  def sqrt(a: StoragePtr, b: StoragePtr, idx_a: int, idx_b: int, size: int): ...
-  @staticmethod
-  def exp(a: StoragePtr, b: StoragePtr, idx_a: int, idx_b: int, size: int): ...
-  @staticmethod
-  def log(a: StoragePtr, b: StoragePtr, idx_a: int, idx_b: int, size: int): ...
-  @staticmethod
-  def add(a: StoragePtr, b: StoragePtr, c: StoragePtr, idx_a: int, idx_b: int, idx_c: int, size: int): ...
-  @staticmethod
-  def sub(a: StoragePtr, b: StoragePtr, c: StoragePtr, idx_a: int, idx_b: int, idx_c: int, size: int): ...
-  @staticmethod
-  def mul(a: StoragePtr, b: StoragePtr, c: StoragePtr, idx_a: int, idx_b: int, idx_c: int, size: int): ...
-  @staticmethod
-  def div(a: StoragePtr, b: StoragePtr, c: StoragePtr, idx_a: int, idx_b: int, idx_c: int, size: int): ...
-  @staticmethod
-  def sum(a: StoragePtr, b: StoragePtr, idx_a: int, idx_b: int, size: int): ...
-  @staticmethod
-  def max(a: StoragePtr, b: StoragePtr, idx_a: int, idx_b: int, size: int): ...
+  def from_vec(vec: Union[list[float], np.ndarray], device: str) -> StorageView: ...
+  def len(self) -> 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 neg(self) -> StorageView: ...
+  def sqrt(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: ...

rust-toolchain.toml

@@ -1,2 +1,2 @@
 [toolchain]
-channel = "nightly-2024-08-15"
+channel = "nightly"