wehub-resource-sync
48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
# Copyright (c) 2025 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 is_custom_device
|
|
|
|
import paddle
|
|
from paddle.base import core
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device()),
|
|
"core is not compiled with CUDA",
|
|
)
|
|
class TestRestrictNonzero(unittest.TestCase):
|
|
def test_restrict_nonzero(self):
|
|
# test dynamic
|
|
paddle.disable_static()
|
|
x = paddle.to_tensor(
|
|
[[-1, 2, 3, -1], [0, 1, 2, -1], [0, -1, 1, -1]]
|
|
).flatten()
|
|
num_tokens_per_expert_list = [2, 2, 2, 1]
|
|
ref_out = (x + (x == -1).cast('int64') * 256).argsort()[:7]
|
|
out = paddle.concat(
|
|
[
|
|
paddle.tensor.search._restrict_nonzero(x == i, total_true_num)
|
|
for i, total_true_num in enumerate(num_tokens_per_expert_list)
|
|
]
|
|
).flatten()
|
|
np.testing.assert_equal(actual=out.numpy(), desired=ref_out.numpy())
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|