com.microsoft - GatherND

GatherND - 1

Version

• name: GatherND (GitHub)

• domain: com.microsoft

• since_version: 1

• function:

• support_level: SupportType.COMMON

• shape inference: True

This version of the operator has been available since version 1 of domain com.microsoft.

Summary

Inputs

• data (heterogeneous) - T:

• indices (heterogeneous) - Tind:

Outputs

• output (heterogeneous) - T:

Type Constraints

• T in ( tensor(bool), tensor(complex128), tensor(complex64), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8) ): Constrain input and output types to any tensor type.

• Tind in ( tensor(int32), tensor(int64) ): Constrain indice type to int32 or int64

Examples

_int32

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GatherND",
    inputs=["data", "indices"],
    outputs=["output"],
)

data = np.array([[0, 1], [2, 3]], dtype=np.int32)
indices = np.array([[0, 0], [1, 1]], dtype=np.int64)
output = gather_nd_impl(data, indices, 0)
expected_output = np.array([0, 3], dtype=np.int32)
assert np.array_equal(output, expected_output)
expect(
    node,
    inputs=[data, indices],
    outputs=[output],
    name="test_gathernd_example_int32",
)

_float32

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GatherND",
    inputs=["data", "indices"],
    outputs=["output"],
)

data = np.array([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dtype=np.float32)
indices = np.array([[[0, 1]], [[1, 0]]], dtype=np.int64)
output = gather_nd_impl(data, indices, 0)
expected_output = np.array([[[2, 3]], [[4, 5]]], dtype=np.float32)
assert np.array_equal(output, expected_output)
expect(
    node,
    inputs=[data, indices],
    outputs=[output],
    name="test_gathernd_example_float32",
)

_int32_batchdim_1

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GatherND",
    inputs=["data", "indices"],
    outputs=["output"],
    batch_dims=1,
)

data = np.array([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dtype=np.int32)
indices = np.array([[1], [0]], dtype=np.int64)
output = gather_nd_impl(data, indices, 1)
expected_output = np.array([[2, 3], [4, 5]], dtype=np.int32)
assert np.array_equal(output, expected_output)
expect(
    node,
    inputs=[data, indices],
    outputs=[output],
    name="test_gathernd_example_int32_batch_dim1",
)