[BugFix][Relay] fix scatter_nd type relation

[JR4er]

import tvm
import numpy as np
from tvm import relay


def should_pass():

    data = relay.const(np.zeros([1,1,10], 'float32'))
    indices = relay.const(np.ones([2,1,1,1], 'int64'))
    updates = relay.const(np.random.rand(1,1,1,10).astype('float32'))

    scatter_nd = relay.scatter_nd(data, indices, updates)

    mod = tvm.IRModule.from_expr(scatter_nd)
    mod = relay.transform.InferType()(mod)


def should_fail():
    data = relay.const(np.zeros([1,1,10], 'float32'))
    indices = relay.const(np.ones([2,1,1], 'int64'))
    updates = relay.const(np.random.rand(1,1,5).astype('float32'))

    scatter_nd = relay.scatter_nd(data, indices, updates)

    mod = tvm.IRModule.from_expr(scatter_nd)
    mod = relay.transform.InferType()(mod)

case should_pass will report a error

Check failed: (0 <= i && i < p->size_) is false: IndexError: indexing 3 on an array of size 3

case should_fail will pass incorrectly. however FoldConstant will find the error:

AssertionError: Dimension of updates[2] (5) must equal dimension of out_shape[2] (10).

It looks like #7927 changes the implementation of relay.scatter_nd into onnx-like. but a change from data into updates is missing.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @shingjan _{See #10317 for details}

_{Generated by tvm-bot}