[Relay] Support dynamic indices size in gather_nd and scatter_nd

[masahi]

Added shape func for gather_nd op to support dynamic indices size. scatter_nd also works with dynamic indices after dropping some assertions check on shapes.

please review @mbrookhart @jwfromm @comaniac @kevinthesun

One difficulty with gather_nd shape func was that the rank of the output tensor depends on the runtime shape of indices tensor.
At the line

tvm/python/tvm/relay/op/_transform.py

Line 1089 in 55a4430

out_shape = output_tensor((kdim + ndim - (mdim + batch_dims),), "int64")

,
if instead I have something like

mdim = indices_shape[0]  # indices_shape is a runtime shape represented as te::Tensor 
out_shape = output_tensor((kdim + ndim - (mdim + batch_dims),), "int64")

different errors are raised depending on the context. For example, an error occurs in the shape func for squeeze that follows gather_dim, at

tvm/python/tvm/relay/op/_transform.py

Line 772 in 55a4430

for i in range(inputs[0].shape[0].value):

. That line is a python range for loop that expects the rank of a shape tensor, inputs[0].shape[0].value, to be a compile time fixed integer. But if the rank of the gather_dim output tensor is calculated like above, the rank would not be a compile time integer (because it depends on the runtime value of shape).

My workaround for this problem is based on the fact that the first axis of indices tensor, whose runtime size is required to compute the output rank, is actually a compile time constant as asserted at

tvm/src/relay/op/tensor/transform.cc

Lines 3347 to 3348 in 55a4430

	const IntImmNode* mdim = indices->shape[0].as<IntImmNode>();
	ICHECK(mdim) << "GatherND needs a static shape for the first axis of indices, got "

. So I simply record this fixed value in the attribute dict and look it up inside shape func. To do this, I had to add a new attribute to gather_dim, which is a bit ad hoc but necessary.

[masahi]

pinging for reviews.