From dc331231b7270e84cf54b663b75fd461e68208be Mon Sep 17 00:00:00 2001 From: Aganlengzi Date: Fri, 19 Aug 2022 09:48:51 +0800 Subject: [PATCH] [CustomDevice] support scalar (#45244) --- paddle/phi/api/lib/scalar.cc | 6 ++++++ .../fluid/tests/custom_runtime/test_custom_cpu_plugin.py | 8 ++++++++ 2 files changed, 14 insertions(+) diff --git a/paddle/phi/api/lib/scalar.cc b/paddle/phi/api/lib/scalar.cc index c31338de09f..09b15f629d8 100644 --- a/paddle/phi/api/lib/scalar.cc +++ b/paddle/phi/api/lib/scalar.cc @@ -37,6 +37,12 @@ ScalarBase::ScalarBase(const Tensor& tensor_in) GetDataFromTensor(dst_tensor); } else if (tensor_in_place == phi::AllocationType::CPU) { GetDataFromTensor(tensor_in); +#ifdef PADDLE_WITH_CUSTOM_DEVICE + } else if (tensor_in_place == phi::AllocationType::CUSTOM) { + Tensor dst_tensor; + copy(tensor_in, phi::CPUPlace(), true, &dst_tensor); + GetDataFromTensor(dst_tensor); +#endif } else { PADDLE_THROW(phi::errors::Unimplemented( "Now, it is not supported to construct Scalar using tensor that its " diff --git a/python/paddle/fluid/tests/custom_runtime/test_custom_cpu_plugin.py b/python/paddle/fluid/tests/custom_runtime/test_custom_cpu_plugin.py index bf1effe2191..07e22516040 100644 --- a/python/paddle/fluid/tests/custom_runtime/test_custom_cpu_plugin.py +++ b/python/paddle/fluid/tests/custom_runtime/test_custom_cpu_plugin.py @@ -41,6 +41,7 @@ class TestCustomCPUPlugin(unittest.TestCase): self._test_eager_backward_api() self._test_eager_copy_to() self._test_fallback_kernel() + self._test_scalar() self._test_custom_device_dataloader() self._test_custom_device_mnist() @@ -170,6 +171,13 @@ class TestCustomCPUPlugin(unittest.TestCase): z = paddle.add(x, y) np.testing.assert_array_equal(z, r) + def _test_scalar(self): + import paddle + data_1 = paddle.to_tensor([[[[1.0, 4.0, 5.0, 7.0], [3.0, 4.0, 5.0, + 6.0]]]]) + k_t = paddle.to_tensor([3], dtype="int32") + value_1, indices_1 = paddle.topk(data_1, k=k_t) + def tearDown(self): del os.environ['CUSTOM_DEVICE_ROOT'] -- GitLab