diff --git a/src/candle/__init__.py b/src/candle/__init__.py index 529c0cff..3dc7c02c 100644 --- a/src/candle/__init__.py +++ b/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 diff --git a/src/candle/_creation.py b/src/candle/_creation.py index 133293d5..fa5a6f0e 100644 --- a/src/candle/_creation.py +++ b/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 diff --git a/src/candle/_storage.py b/src/candle/_storage.py index cb54dfd6..7ecea516 100644 --- a/src/candle/_storage.py +++ b/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) diff --git a/tests/test_torch_compat.py b/tests/test_torch_compat.py index a079418c..3b9a5a5a 100644 --- a/tests/test_torch_compat.py +++ b/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("""\