From 3f64c96655289a1874b968d6bd9287eeb99a2cec Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Tue, 24 Jan 2023 08:09:30 +0000 Subject: [PATCH] `asarray`: Add support for NumPy scalars (#90914) Follow up from: Quansight-Labs/numpy_pytorch_interop#3 This PR adds support for NumPy scalars for `torch.asarray`. **Before:** treats the scalar as an object that implements the buffer protocol. Thus, interprets the data as the default data type (`float32`) ```python >>> torch.asarray(numpy.float64(0.5)) tensor([0.0000, 1.7500]) ``` **After:** identifies the NumPy scalar, and does the "right" thing. i.e. creates a 0-dimensional tensor from the NumPy array that doesn't share its memory ```python >>> torch.asarray(numpy.float64(0.5)) tensor(0.5000, dtype=torch.float64) ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/90914 Approved by: https://github.com/lezcano, https://github.com/mruberry --- test/test_tensor_creation_ops.py | 12 +++++++++++ torch/_torch_docs.py | 14 +++++++++--- torch/csrc/utils/tensor_new.cpp | 37 ++++++++++++++++++++++++++++---- 3 files changed, 56 insertions(+), 7 deletions(-) diff --git a/test/test_tensor_creation_ops.py b/test/test_tensor_creation_ops.py index 67accdecb1742b..13e6f399d7a643 100644 --- a/test/test_tensor_creation_ops.py +++ b/test/test_tensor_creation_ops.py @@ -3936,6 +3936,18 @@ def test_astensor_consistency(self, device): t = torch.asarray(e) self.assertEqual(t, original) + @onlyCPU + def test_numpy_scalars(self, device): + scalar = np.float64(0.5) + + with self.assertRaisesRegex(RuntimeError, "can't alias NumPy scalars."): + torch.asarray(scalar, copy=False) + + tensor = torch.asarray(scalar) + self.assertEqual(tensor.dim(), 0) + self.assertEqual(tensor.item(), scalar.item()) + self.assertEqual(tensor.dtype, torch.float64) + instantiate_device_type_tests(TestTensorCreation, globals()) instantiate_device_type_tests(TestRandomTensorCreation, globals()) instantiate_device_type_tests(TestLikeTensorCreation, globals()) diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py index 12ed8e037e95d2..18c3e56358c4c7 100644 --- a/torch/_torch_docs.py +++ b/torch/_torch_docs.py @@ -1230,7 +1230,7 @@ def merge_dicts(*dicts): :attr:`obj` can be one of: 1. a tensor -2. a NumPy array +2. a NumPy array or a NumPy scalar 3. a DLPack capsule 4. an object that implements Python's buffer protocol 5. a scalar @@ -1245,14 +1245,18 @@ def merge_dicts(*dicts): is ``True`` then the returned tensor will require a gradient, and if :attr:`obj` is also a tensor with an autograd history then the returned tensor will have the same history. -When :attr:`obj` is not a tensor, NumPy Array, or DLPack capsule but implements Python's +When :attr:`obj` is not a tensor, NumPy array, or DLPack capsule but implements Python's buffer protocol then the buffer is interpreted as an array of bytes grouped according to the size of the datatype passed to the :attr:`dtype` keyword argument. (If no datatype is passed then the default floating point datatype is used, instead.) The returned tensor will have the specified datatype (or default floating point datatype if none is specified) and, by default, be on the CPU device and share memory with the buffer. -When :attr:`obj` is none of the above but a scalar or sequence of scalars then the +When :attr:`obj` is a NumPy scalar, the returned tensor will be a 0-dimensional tensor on +the CPU and that doesn't share its memory (i.e. ``copy=True``). By default datatype will +be the PyTorch datatype corresponding to the NumPy's scalar's datatype. + +When :attr:`obj` is none of the above but a scalar, or a sequence of scalars then the returned tensor will, by default, infer its datatype from the scalar values, be on the CPU device, and not share its memory. @@ -1320,6 +1324,10 @@ def merge_dicts(*dicts): >>> t2 = torch.asarray(array, dtype=torch.float32) >>> array.__array_interface__['data'][0] == t1.data_ptr() False + + >>> scalar = numpy.float64(0.5) + >>> torch.asarray(scalar) + tensor(0.5000, dtype=torch.float64) """, ) diff --git a/torch/csrc/utils/tensor_new.cpp b/torch/csrc/utils/tensor_new.cpp index 6121c4c43eed1d..4d0abf864b21b0 100644 --- a/torch/csrc/utils/tensor_new.cpp +++ b/torch/csrc/utils/tensor_new.cpp @@ -1603,10 +1603,39 @@ Tensor asarray( } #ifdef USE_NUMPY - // Check whether 'obj' is a NumPy Array - if (is_numpy_available() && PyArray_Check(obj)) { - tensor = tensor_from_numpy(obj, /*warn_if_not_writeable=*/false); - should_warn_numpy_not_writable = !PyArray_ISWRITEABLE((PyArrayObject*)obj); + if (is_numpy_available()) { + // Check whether 'obj' is a NumPy Array or Scalar. + bool is_numpy_array = PyArray_Check(obj); + bool is_numpy_scalar = PyArray_CheckScalar(obj); + + if (is_numpy_array || is_numpy_scalar) { + THPObjectPtr ptr; + auto arr = obj; + + if (is_numpy_scalar) { + TORCH_CHECK( + !force_alias, + "can't alias NumPy scalars. ", + "Either remove copy=False or transform it in a ndarray. ") + + ptr = PyArray_FromScalar(obj, nullptr); + arr = ptr.get(); + } + + tensor = tensor_from_numpy(arr, /*warn_if_not_writeable=*/false); + should_warn_numpy_not_writable = + !PyArray_ISWRITEABLE((PyArrayObject*)arr); + + if (is_numpy_scalar) { + // Uses a newly cloned storage, instead of the shared one. + // The THPObjectPtr will delete the previous storage in the + // end of the previous scope. + tensor = tensor.clone(); + + // No need to clone again, later. + force_copy = false; + } + } } #endif