Spaces:
Running
Running
File size: 2,809 Bytes
dc2106c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
# Copyright (c) ONNX Project Contributors
#
# SPDX-License-Identifier: Apache-2.0
import numpy as np
import onnx
from onnx.backend.test.case.base import Base
from onnx.backend.test.case.node import expect
class Compress(Base):
@staticmethod
def export_compress_0() -> None:
node = onnx.helper.make_node(
"Compress",
inputs=["input", "condition"],
outputs=["output"],
axis=0,
)
input = np.array([[1, 2], [3, 4], [5, 6]]).astype(np.float32)
condition = np.array([0, 1, 1])
output = np.compress(condition, input, axis=0)
# print(output)
# [[ 3. 4.]
# [ 5. 6.]]
expect(
node,
inputs=[input, condition.astype(bool)],
outputs=[output],
name="test_compress_0",
)
@staticmethod
def export_compress_1() -> None:
node = onnx.helper.make_node(
"Compress",
inputs=["input", "condition"],
outputs=["output"],
axis=1,
)
input = np.array([[1, 2], [3, 4], [5, 6]]).astype(np.float32)
condition = np.array([0, 1])
output = np.compress(condition, input, axis=1)
# print(output)
# [[ 2.]
# [ 4.]
# [ 6.]]
expect(
node,
inputs=[input, condition.astype(bool)],
outputs=[output],
name="test_compress_1",
)
@staticmethod
def export_compress_default_axis() -> None:
node = onnx.helper.make_node(
"Compress",
inputs=["input", "condition"],
outputs=["output"],
)
input = np.array([[1, 2], [3, 4], [5, 6]]).astype(np.float32)
condition = np.array([0, 1, 0, 0, 1])
output = np.compress(condition, input)
# print(output)
# [ 2., 5.]
expect(
node,
inputs=[input, condition.astype(bool)],
outputs=[output],
name="test_compress_default_axis",
)
@staticmethod
def export_compress_negative_axis() -> None:
node = onnx.helper.make_node(
"Compress",
inputs=["input", "condition"],
outputs=["output"],
axis=-1,
)
input = np.array([[1, 2], [3, 4], [5, 6]]).astype(np.float32)
condition = np.array([0, 1])
output = np.compress(condition, input, axis=-1)
# print(output)
# [[ 2.]
# [ 4.]
# [ 6.]]
expect(
node,
inputs=[input, condition.astype(bool)],
outputs=[output],
name="test_compress_negative_axis",
)
|