wehub-resource-sync
208 lines
6.8 KiB
Python
208 lines
6.8 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from op_test import get_device_place, get_places, is_custom_device
|
|
|
|
import paddle
|
|
|
|
|
|
def generate_data(shape, dtype):
|
|
"""
|
|
Generate `data` and `mask` with the same shape and dtype.
|
|
"""
|
|
_mask = np.random.randint(0, 2, shape)
|
|
if np.sum(_mask) == 0:
|
|
_mask.flat[0] = 1
|
|
mask = (np.random.randint(-100, 100, shape) * _mask).astype(dtype)
|
|
data = np.random.randint(-100, 100, shape).astype(dtype)
|
|
return data, mask
|
|
|
|
|
|
class TestMaskAs(unittest.TestCase):
|
|
def setUp(self):
|
|
self.init_format()
|
|
self.places = get_places()
|
|
|
|
def init_format(self):
|
|
self.format = None
|
|
|
|
def check(self, shape, dtype, place, check_grad=True):
|
|
paddle.disable_static()
|
|
dense_data_np, dense_mask_np = generate_data(shape, dtype)
|
|
|
|
dense_data_pd = paddle.to_tensor(
|
|
dense_data_np, dtype=dtype, place=place
|
|
)
|
|
dense_data_pd.stop_gradient = False
|
|
|
|
if self.format == 'coo':
|
|
sparse_mask_pd = paddle.to_tensor(
|
|
dense_mask_np, dtype=dtype, place=place
|
|
).to_sparse_coo(len(shape))
|
|
else:
|
|
sparse_mask_pd = paddle.to_tensor(
|
|
dense_mask_np, dtype=dtype, place=place
|
|
).to_sparse_csr()
|
|
|
|
sparse_out_pd = paddle.sparse.mask_as(dense_data_pd, sparse_mask_pd)
|
|
|
|
# compare the tensor from sparse->dense with reference numpy data
|
|
# the result only keeps the values where mask not zero, like:
|
|
# dense_data_np
|
|
# [[ 38. 15. 76.]
|
|
# [-98. -75. 10.]
|
|
# [-52. 49. -48.]]
|
|
# dense_mask_np
|
|
# [[-70. 0. 0.]
|
|
# [-50. 34. 60.]
|
|
# [-34. 0. -18.]]
|
|
# dense_data_np_ref
|
|
# [[ 38. 0. 0.]
|
|
# [-98. -75. 10.]
|
|
# [-52. 0. -48.]]
|
|
dense_data_np_ref = dense_data_np * (dense_mask_np != 0)
|
|
np.testing.assert_allclose(
|
|
sparse_out_pd.to_dense().numpy(), dense_data_np_ref
|
|
)
|
|
|
|
if check_grad:
|
|
# with sparse_out_pd backward, we get the grad from dense_data_pd
|
|
sparse_out_pd.backward()
|
|
dense_data_grad = dense_data_pd.grad
|
|
|
|
self.assertEqual(
|
|
list(dense_data_grad.shape), list(dense_data_pd.shape)
|
|
)
|
|
self.assertEqual(dense_data_grad.dtype, dense_data_pd.dtype)
|
|
|
|
# make a dense data to compare the grad from sparse_out_pd
|
|
grad_ref = np.ones_like(dense_mask_np) * (dense_mask_np != 0)
|
|
|
|
np.testing.assert_allclose(
|
|
dense_data_pd.grad.numpy(),
|
|
grad_ref,
|
|
)
|
|
|
|
def check_with_dtypes(self, shape):
|
|
for place in self.places:
|
|
self.check(shape, 'float32', place)
|
|
self.check(shape, 'float64', place)
|
|
self.check(shape, 'int32', place)
|
|
self.check(shape, 'int64', place)
|
|
self.check(shape, 'complex64', place)
|
|
self.check(shape, 'complex128', place)
|
|
|
|
# `int8`` not registered in `FullLikeCooKernel`, so skip check_grad
|
|
self.check(shape, 'int8', place, check_grad=False)
|
|
|
|
# `int16` not registered in `multiply`, so skip check_grad
|
|
self.check(shape, 'int16', place, check_grad=False)
|
|
|
|
if paddle.is_compiled_with_cuda() or is_custom_device():
|
|
place = get_device_place()
|
|
self.check(shape, 'float16', place)
|
|
|
|
def test_tensor_sparse_mask(self):
|
|
"""
|
|
Test the tensor method `sparse_mask` against the paddle.sparse.mask_as API
|
|
"""
|
|
# Test for 1D, 2D, 3D, 4D tensors
|
|
shapes = [(5,), (5, 3), (5, 3, 4), (5, 3, 4, 2)]
|
|
for shape in shapes:
|
|
for dtype in [
|
|
'float32',
|
|
'float64',
|
|
'int32',
|
|
'int64',
|
|
'complex64',
|
|
'complex128',
|
|
]:
|
|
for place in self.places:
|
|
# Generate data
|
|
dense_data_np, dense_mask_np = generate_data(shape, dtype)
|
|
dense_data_pd = paddle.to_tensor(
|
|
dense_data_np, dtype=dtype, place=place
|
|
)
|
|
dense_data_pd.stop_gradient = False
|
|
|
|
# Convert mask to sparse
|
|
sparse_mask_pd = paddle.to_tensor(
|
|
dense_mask_np, dtype=dtype, place=place
|
|
).to_sparse_coo(len(shape))
|
|
|
|
# Use the new tensor method (your API)
|
|
sparse_out_pd = dense_data_pd.sparse_mask(sparse_mask_pd)
|
|
|
|
# Compare with reference (same as original test)
|
|
dense_data_np_ref = dense_data_np * (dense_mask_np != 0)
|
|
np.testing.assert_allclose(
|
|
sparse_out_pd.to_dense().numpy(), dense_data_np_ref
|
|
)
|
|
|
|
# Check gradient (skip int8 and int16)
|
|
if dtype not in ['int8', 'int16']:
|
|
sparse_out_pd.backward()
|
|
dense_data_grad = dense_data_pd.grad
|
|
grad_ref = np.ones_like(dense_mask_np) * (
|
|
dense_mask_np != 0
|
|
)
|
|
np.testing.assert_allclose(
|
|
dense_data_grad.numpy(),
|
|
grad_ref,
|
|
)
|
|
|
|
|
|
class TestMaskAsCoo(TestMaskAs):
|
|
def init_format(self):
|
|
self.format = 'coo'
|
|
|
|
def test_1d(self):
|
|
self.check_with_dtypes((5,))
|
|
|
|
def test_2d(self):
|
|
self.check_with_dtypes((5, 3))
|
|
|
|
def test_3d(self):
|
|
self.check_with_dtypes((5, 3, 4))
|
|
|
|
def test_4d(self):
|
|
self.check_with_dtypes((5, 3, 4, 2))
|
|
|
|
|
|
class TestMaskAsCsr(TestMaskAs):
|
|
def init_format(self):
|
|
self.format = 'csr'
|
|
|
|
def test_2d(self):
|
|
self.check_with_dtypes((5, 3))
|
|
|
|
def test_3d(self):
|
|
self.check_with_dtypes((5, 3, 4))
|
|
|
|
def test_error_dimension(self):
|
|
# error 1d
|
|
with self.assertRaises(ValueError):
|
|
self.check_with_dtypes((5,))
|
|
|
|
# error 4d
|
|
with self.assertRaises(ValueError):
|
|
self.check_with_dtypes((5, 3, 4, 2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|