fix: add torch.frombuffer compatibility

src/candle/__init__.py

@@ -90,7 +90,7 @@ def quantize_per_tensor(*_args, **_kwargs):
 ComplexFloatTensor = Tensor
 ComplexDoubleTensor = Tensor
 Size = tuple
-from ._creation import tensor, zeros, ones, empty, arange, linspace, full, logspace, eye, range, randn, rand, randint, randperm, from_numpy, as_tensor, normal
+from ._creation import tensor, zeros, ones, empty, arange, linspace, full, logspace, eye, range, randn, rand, randint, randperm, from_numpy, frombuffer, as_tensor, normal
 from ._functional import zeros_like
 from ._functional import ones_like, empty_like, full_like, randn_like, rand_like, randint_like
 from ._storage import UntypedStorage, TypedStorage

src/candle/_creation.py

@@ -1,5 +1,6 @@
 import numpy as np
 from ._dtype import float32, int64
+from ._dtype import to_numpy_dtype
 from ._dtype import bool as bool_dtype
 from ._functional import tensor as tensor_dispatch
 from ._functional import zeros as zeros_dispatch
@@ -16,6 +17,8 @@
 from ._functional import randint as randint_dispatch
 from ._functional import randperm as randperm_dispatch
 from ._functional import normal as normal_dispatch
+from ._storage import typed_storage_from_numpy_view
+from ._cython._tensor_impl import cy_make_tensor_from_storage  # pylint: disable=import-error,no-name-in-module
 
 
 def _apply_requires_grad(out, requires_grad):
@@ -166,6 +169,15 @@ def from_numpy(ndarray):
     return tensor_dispatch(ndarray, dtype=dt)
 
 
+def frombuffer(buffer, *, dtype, count=-1, offset=0, requires_grad=False):
+    np_dtype = to_numpy_dtype(dtype)
+    arr = np.frombuffer(buffer, dtype=np_dtype, count=count, offset=offset)
+    storage = typed_storage_from_numpy_view(arr, dtype)
+    stride = tuple(np.array(arr.strides) // arr.itemsize)
+    out = cy_make_tensor_from_storage(storage, arr.shape, stride, 0, False)
+    return _apply_requires_grad(out, requires_grad)
+
+
 def as_tensor(data, dtype=None, device=None):
     from ._tensor import Tensor
 

src/candle/_storage.py

@@ -586,6 +586,14 @@ def typed_storage_from_numpy(arr, dtype, device=None):
     return TypedStorage(untyped, dtype, arr.size, data=arr)
 
 
+def typed_storage_from_numpy_view(arr, dtype, device=None):
+    arr = np.asarray(arr, dtype=to_numpy_dtype(dtype))
+    if not arr.flags.c_contiguous:
+        raise ValueError("expected contiguous numpy view")
+    untyped = _CPUUntypedStorage(arr.view(np.uint8), device=device)
+    return TypedStorage(untyped, dtype, arr.size, data=arr)
+
+
 def empty_cpu_typed_storage(shape, dtype, device=None):
     arr = np.empty(shape, dtype=to_numpy_dtype(dtype))
     return typed_storage_from_numpy(arr, dtype, device=device)

tests/test_torch_compat.py

@@ -305,6 +305,42 @@ def test_from_torch_utils_model_zoo_imports_load_url(self):
         out, _, _ = _run(code, env_extra={"USE_CANDLE": "1"})
         assert "OK" in out
 
+    def test_torch_frombuffer_shares_memory_with_writable_buffer(self):
+        """torch.frombuffer must alias the original writable buffer."""
+        code = textwrap.dedent("""\
+            import torch
+            raw = bytearray([1, 2, 3, 4])
+            t = torch.frombuffer(raw, dtype=torch.uint8)
+            raw[1] = 99
+            assert t.tolist() == [1, 99, 3, 4]
+            print("OK")
+        """)
+        out, _, _ = _run(code, env_extra={"USE_CANDLE": "1"})
+        assert "OK" in out
+
+    def test_torch_frombuffer_accepts_bytearray_with_offset(self):
+        """Torchvision MNIST loaders use torch.frombuffer on raw IDX bytes."""
+        code = textwrap.dedent("""\
+            import torch
+            t = torch.frombuffer(bytearray(range(8)), dtype=torch.uint8, offset=4)
+            assert t.tolist() == [4, 5, 6, 7]
+            print("OK")
+        """)
+        out, _, _ = _run(code, env_extra={"USE_CANDLE": "1"})
+        assert "OK" in out
+
+    def test_torch_frombuffer_supports_multibyte_dtype_and_view(self):
+        """Torchvision MNIST uses frombuffer result with shape reinterpretation."""
+        code = textwrap.dedent("""\
+            import torch
+            raw = bytearray([1, 0, 2, 0, 3, 0, 4, 0])
+            t = torch.frombuffer(raw, dtype=torch.int16)
+            assert t.view(2, 2).tolist() == [[1, 2], [3, 4]]
+            print("OK")
+        """)
+        out, _, _ = _run(code, env_extra={"USE_CANDLE": "1"})
+        assert "OK" in out
+
     def test_torch_zeros(self):
         """Functional: actually create a tensor through the redirected import."""
         code = textwrap.dedent("""\