gathernd fix

Author: wayuanho

tf2onnx/function/gathernd.py

@@ -49,7 +49,7 @@ def _make_gathernd_inner_loop(ctx, params, index, dtype):
     return inner_loop
 
 
-def make_gathernd(ctx, params, indices, output, scope_name, t_params):
+def make_gathernd(ctx, params, indices, output, scope_name, t_params, shapes, dtypes):
     """make GatherNd op."""
     # Tparams output = GatherNd(Tparams params, Tidx indices)
     scope_name = utils.make_name(scope_name)
@@ -131,7 +131,11 @@ def make_gathernd(ctx, params, indices, output, scope_name, t_params):
                                  [output_shape_.output[0]],
                                  attr={"axes": [0], "ends": [-1], "starts": [0]},
                                  dtypes=[TensorProto.INT64])
-    ctx.make_node("Reshape", [gathernd_loop.output[1], output_shape.output[0]], outputs=[output])
+    ctx.make_node("Reshape",
+                  [gathernd_loop.output[1], output_shape.output[0]],
+                  outputs=[output],
+                  shapes=shapes,
+                  dtypes=dtypes)
 
 
 def gathernd_op(ctx, node, name, args):
@@ -143,4 +147,7 @@ def gathernd_op(ctx, node, name, args):
     # same as the attr Tparams
     t_params = ctx.get_dtype(params)
     utils.make_sure(t_params, "Dtype of {} is None".format(indices))
-    make_gathernd(ctx, params, indices, output, name, t_params)
+    shapes = node.output_shapes
+    dtypes = node.output_dtypes
+    ctx.remove_node(node.name)
+    make_gathernd(ctx, params, indices, output, name, t_params, shapes, dtypes)

tf2onnx/function/sparse_softmax_cross_entropy_with_logits.py

@@ -59,6 +59,7 @@ def sparse_softmax_cross_entropy_with_logits_op_by_gathernd(ctx, node, name, arg
         raise ValueError("onehot op: only rank1 is supported")
     logit_name = node.input[0]
     logit_dtype = ctx.get_dtype(logit_name)
+    logit_shape = ctx.get_shape(logit_name)
     utils.make_sure(logit_dtype, "Dtype of {} is None".format(logit_name))
     indices_dtype = ctx.get_dtype(indices_name)
     if indices_dtype != TensorProto.INT64:
@@ -76,11 +77,12 @@ def sparse_softmax_cross_entropy_with_logits_op_by_gathernd(ctx, node, name, arg
     indices_with_id = ctx.make_node("Concat",
                                     [id_unsqueeze.output[0], indices_unsqueeze.output[0]],
                                     attr={"axis": 1})
-    log_softmax = ctx.make_node(op_type="LogSoftmax", inputs=[logit_name], dtypes=[logit_dtype])
+    log_softmax = ctx.make_node(op_type="LogSoftmax",
+                                inputs=[logit_name], dtypes=[logit_dtype], shapes=[logit_shape])
     gathernd_name = utils.make_name("sparse_softmax_gathernd")
     gathernd_output = utils.port_name(gathernd_name)
     make_gathernd(ctx, log_softmax.output[0], indices_with_id.output[0], gathernd_output,
-                  gathernd_name, logit_dtype)
+                  gathernd_name, logit_dtype, [logit_shape], [logit_dtype])
     const_name = utils.make_name("const_negative_one")
     const_negative_one = ctx.make_const(const_name, np.array(-1).astype(utils.map_onnx_to_numpy_type(logit_dtype)))
     mul2 = ctx.make_node(op_type="Mul", inputs=[const_negative_one.output[0], gathernd_output])