mirror of
https://github.com/luxfi/onnx-source-snapshot.git
synced 2026-10-30 15:23:38 +00:00
175 lines
5.1 KiB
Python
175 lines
5.1 KiB
Python
# Copyright (c) ONNX Project Contributors
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
import onnx
|
|
from onnx.backend.test.case.base import Base
|
|
from onnx.backend.test.case.node import expect
|
|
|
|
|
|
class TensorScatter(Base):
|
|
@staticmethod
|
|
def export_tensorscatter() -> None:
|
|
node = onnx.helper.make_node(
|
|
"TensorScatter",
|
|
inputs=["past_cache", "update", "write_indices"],
|
|
outputs=["present_cache"],
|
|
mode="linear",
|
|
)
|
|
past_cache = np.array(
|
|
[
|
|
[[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
|
|
[[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
update = np.array(
|
|
[
|
|
[[[5, 5, 5, 5, 5]]],
|
|
[[[1, 1, 1, 1, 1]]],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
write_indices = np.array([1, 2], dtype=np.int64)
|
|
present_cache = np.array(
|
|
[
|
|
[[[1, 2, 3, 4, 5], [5, 5, 5, 5, 5], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
|
|
[[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [1, 1, 1, 1, 1], [4, 3, 2, 1, 0]]],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
expect(
|
|
node,
|
|
inputs=[past_cache, update, write_indices],
|
|
outputs=[present_cache],
|
|
name="test_tensorscatter",
|
|
)
|
|
|
|
@staticmethod
|
|
def export_tensorscatter_circular() -> None:
|
|
node = onnx.helper.make_node(
|
|
"TensorScatter",
|
|
inputs=["past_cache", "update", "write_indices"],
|
|
outputs=["present_cache"],
|
|
mode="circular",
|
|
)
|
|
past_cache = np.array(
|
|
[
|
|
[[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
|
|
[[[1, 2, 3, 4, 5], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [4, 3, 2, 1, 0]]],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
update = np.array(
|
|
[
|
|
[
|
|
[
|
|
[5, 5, 5, 5, 5],
|
|
[6, 6, 6, 6, 6],
|
|
]
|
|
],
|
|
[
|
|
[
|
|
[1, 1, 1, 1, 1],
|
|
[2, 2, 2, 2, 2],
|
|
]
|
|
],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
write_indices = np.array([1, 3], dtype=np.int64)
|
|
present_cache = np.array(
|
|
[
|
|
[[[1, 2, 3, 4, 5], [5, 5, 5, 5, 5], [6, 6, 6, 6, 6], [4, 3, 2, 1, 0]]],
|
|
[[[2, 2, 2, 2, 2], [5, 6, 7, 8, 9], [8, 7, 6, 5, 4], [1, 1, 1, 1, 1]]],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
expect(
|
|
node,
|
|
inputs=[past_cache, update, write_indices],
|
|
outputs=[present_cache],
|
|
name="test_tensorscatter_circular",
|
|
)
|
|
|
|
@staticmethod
|
|
def export_tensorscatter_3d() -> None:
|
|
node = onnx.helper.make_node(
|
|
"TensorScatter",
|
|
inputs=["past_cache", "update", "write_indices"],
|
|
outputs=["present_cache"],
|
|
)
|
|
past_cache = np.array(
|
|
[
|
|
[
|
|
[1, 2, 3, 4, 5],
|
|
[5, 6, 7, 8, 9],
|
|
[8, 7, 6, 5, 4],
|
|
[5, 4, 3, 2, 1],
|
|
],
|
|
[
|
|
[1, 2, 3, 4, 5],
|
|
[5, 6, 7, 8, 9],
|
|
[8, 7, 6, 5, 4],
|
|
[5, 4, 3, 2, 1],
|
|
],
|
|
[
|
|
[1, 2, 3, 4, 5],
|
|
[5, 6, 7, 8, 9],
|
|
[8, 7, 6, 5, 4],
|
|
[5, 4, 3, 2, 1],
|
|
],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
update = np.array(
|
|
[
|
|
[
|
|
[4, 4, 4, 4, 4],
|
|
[5, 5, 5, 5, 5],
|
|
],
|
|
[
|
|
[6, 6, 6, 6, 6],
|
|
[7, 7, 7, 7, 7],
|
|
],
|
|
[
|
|
[2, 2, 2, 2, 2],
|
|
[3, 3, 3, 3, 3],
|
|
],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
write_indices = np.array([1, 2, 0], dtype=np.int64)
|
|
present_cache = np.array(
|
|
[
|
|
[
|
|
[1, 2, 3, 4, 5],
|
|
[4, 4, 4, 4, 4],
|
|
[5, 5, 5, 5, 5],
|
|
[5, 4, 3, 2, 1],
|
|
],
|
|
[
|
|
[1, 2, 3, 4, 5],
|
|
[5, 6, 7, 8, 9],
|
|
[6, 6, 6, 6, 6],
|
|
[7, 7, 7, 7, 7],
|
|
],
|
|
[
|
|
[2, 2, 2, 2, 2],
|
|
[3, 3, 3, 3, 3],
|
|
[8, 7, 6, 5, 4],
|
|
[5, 4, 3, 2, 1],
|
|
],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
expect(
|
|
node,
|
|
inputs=[past_cache, update, write_indices],
|
|
outputs=[present_cache],
|
|
name="test_tensorscatter_3d",
|
|
)
|