Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def select_op9(ctx, node, name, args):
# T output = Select(bool condition, T x, T y)
# T1 output = Where(bool condition, T1 x, T1 y)
# NOTE: condition can be 1-dimension in tensorflow, while in onnx,
# it should be broadcastable with other two inputs
cond_shape = ctx.get_shape(node.input[0])
make_sure(cond_shape is not None, "shape of {} is None".format(node.input[0]))
input_shape = ctx.get_shape(node.input[1])
if input_shape is None:
input_shape = ctx.get_shape(node.input[2])
make_sure(input_shape is not None, "input shape of {} is None".format(node.name))
input_rank = len(input_shape)
# if cond shape is 1-dimensional while input has higher rank, need to be reshaped to broadcast
if len(cond_shape) == 1 and input_rank > 1:
broadcast_shape = [cond_shape[0]] + [1] * (input_rank - 1)
shape_const = ctx.make_const(utils.make_name(name), np.array(broadcast_shape, dtype=np.int64))
reshape = ctx.make_node("Reshape", [node.input[0], shape_const.output[0]])
ctx.replace_input(node, node.input[0], reshape.output[0])
def create_onnx_random_uniform_op(g, tmax, tmin, ru_op, output, to_delete):
dtype = g.get_dtype(output.output[0])
op_name = utils.make_name("RandomUniform")
shape_node = ru_op.inputs[0]
shape = g.get_shape(output.output[0])
if shape_node.is_const():
# if the tensorflow input (aka the shape) is const we can use the RandomUniform op
new_node = g.make_node("RandomUniform", [], name=op_name,
attr={"low": tmin, "high": tmax, "dtype": dtype, "shape": shape},
shapes=[shape], dtypes=[dtype])
else:
if shape_node.type == "Shape":
# if shape is dynamic - in tensorflow shape comes as tensor VALUE,
# in onnx RandomUniformLike finds takes the shape from the tensor itself.
# In many cases there is a shape op in tensorflow before RandomUniform and
# to make that work for onnx we just need to remove the shape op.
new_node = g.make_node("RandomUniformLike", inputs=[shape_node.input[0]], name=op_name,
attr={"low": tmin, "high": tmax, "dtype": dtype},
shapes=[shape], dtypes=[dtype])
attr={"to": TensorProto.INT64})
# cond
# Use initializer here since Constant OP before opset 9 does not support bool type
cond_name = "{}_cond".format(base_name)
ctx.make_const(cond_name, np.ones((), dtype=bool))
# body
g = ctx.create_new_graph_with_same_config()
g.parent_graph = ctx
g.add_graph_input("i", TensorProto.INT64, [])
g.add_graph_input("cond", TensorProto.BOOL, [])
g.add_graph_input("prev", dtype, [])
g.make_node("Identity", ["cond"], outputs=["cond_out"])
g.make_node("Add", ["prev", delta], outputs=["current"], name=utils.make_name("add"))
g.make_node("Identity", ["prev"], outputs=["range"])
g.add_graph_output("cond_out", TensorProto.BOOL, [])
g.add_graph_output("current", dtype, [])
g.add_graph_output("range", dtype, [])
# loop
loop_inputs = [trip_count_node.output[0], cond_name, start]
loop_node = ctx.make_node("Loop", loop_inputs, output_count=2, op_name_scope=base_name, name="loop")
loop_node.set_body_graph_as_attr("body", g)
ctx.make_node("Identity", [loop_node.output[1]], name=base_name, shapes=[shape], dtypes=[dtype], outputs=[output])
def make_range_non_const(ctx, start, limit, delta, output, scope_name, shape, dtype):
"""make Range subgraph."""
# T range = Range(T start, T limit, T delta)
# V v_final_and_scan_outputs = Loop(int64 M, B cond, V v_initial)
base_name = utils.make_name(scope_name)
# trip_count
diff_node = ctx.make_node("Sub",
[limit, start],
op_name_scope=base_name,
name=utils.make_name("diff"))
diff_output = diff_node.output[0]
delta_cast = delta
if dtype in [TensorProto.INT32, TensorProto.INT64]:
cast_node = ctx.make_node("Cast", [diff_output], op_name_scope=base_name,
name="cast_diff", attr={"to": TensorProto.FLOAT})
diff_output = cast_node.output[0]
cast_node = ctx.make_node("Cast", [delta], op_name_scope=base_name, name="cast_delta",
attr={"to": TensorProto.FLOAT})
indices_name = indices_cast.output[0]
indices_size = ctx.make_node("Size", [indices_name])
indices_unsqueeze = ctx.make_node("Unsqueeze", [indices_name], attr={"axes": [1]})
zero_const = ctx.make_const(utils.make_name("zero"), np.array(0, dtype=np.int64))
one_const = ctx.make_const(utils.make_name("one"), np.array(1, dtype=np.int64))
id_name = utils.make_name("sparse_softmax_id")
id_output = utils.port_name(id_name)
controlflow.make_range(ctx, zero_const.output[0], indices_size.output[0], one_const.output[0],
id_output, id_name, shape=[-1], dtype=TensorProto.INT64)
id_unsqueeze = ctx.make_node("Unsqueeze", [id_output], attr={"axes": [1]})
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], shapes=[logit_shape])
gathernd_name = utils.make_name("sparse_softmax_gathernd")
gathernd_output = utils.port_name(gathernd_name)
tensor.make_gathernd(ctx, log_softmax.output[0], indices_with_id.output[0], gathernd_output,
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])
shapes = node.output_shapes
dtypes = node.output_dtypes
ctx.remove_node(node.name)
ctx.make_node(op_type="Squeeze",
inputs=[mul2.output[0]], outputs=[node.output[0]],
attr={"axes": [1]}, shapes=[shapes[0]], dtypes=[dtypes[0]])
# mask end
new_end_mask = np.array(new_end_mask, dtype=np_dtype)
end_output = end.output[0]
if not np.all(new_end_mask == min_size):
if end.is_const():
end = ctx.make_const(
utils.make_name("end_masked"),
np.maximum(end.get_tensor_value(as_list=False), new_end_mask)
)
end_output = end.output[0]
else:
end_mask_const = ctx.make_const(
utils.make_name("end_mask"),
np.array(new_end_mask, dtype=np_dtype)
)
end_output = utils.make_name("{}__end".format(node.name))
math.make_min_or_max_op(ctx, "Max", [end.output[0], end_mask_const.output[0]], [end_output])
# mask strides for shrink
shrink_strided_mask = np.array(shrink_strided_mask, dtype=np_dtype)
strides_output = strides.output[0]
if not np.all(shrink_strided_mask == min_size):
if strides.is_const():
strides = ctx.make_const(
utils.make_name("strides_masked"),
np.maximum(strides.get_tensor_value(as_list=False), shrink_strided_mask)
)
strides_output = strides.output[0]
else:
shrink_strided_mask_const = ctx.make_const(
utils.make_name("strides_mask"),
np.array(shrink_strided_mask, dtype=np_dtype)
)
label = label.output[0]
label_dtype = ctx.get_dtype(label)
logit_dtype = ctx.get_dtype(logit)
utils.make_sure(label_dtype == logit_dtype, "the following logic only works on same dtype of label and logit")
# when label is onehot, logic "tf.multiply(-1, tf.reduce_sum(tf.multiply(label, log_softmax), axis=1))" is equal to
# "-log(q_i)" where i is the selected index specified by label, q_i = logic_i/sum, the detail process is as follows:
# logit_exp=exp(logit) >> sum = tf.reduce_sum(logit_exp, axis = -1), masked_sum = reduce_sum(mul(logit_exp, mul))
# >> -log(masked_sum/sum)
logit_exp = ctx.make_node(op_type="Exp", inputs=[logit]).output[0]
logit_exp_sum = ctx.make_node(op_type="ReduceSum", inputs=[logit_exp], attr={"axes": [-1], "keepdims": 0}).output[0]
masked = ctx.make_node(op_type="Mul", inputs=[label, logit_exp]).output[0]
masked_sum = ctx.make_node(op_type="ReduceSum", inputs=[masked], attr={"axes": [-1], "keepdims": 0}).output[0]
probability = ctx.make_node(op_type="Div", inputs=[masked_sum, logit_exp_sum]).output[0]
log_prob = ctx.make_node(op_type="Log", inputs=[probability]).output[0]
const_negative_one = ctx.make_const(name=utils.make_name("const_negative_one"),
np_val=np.array(-1).astype(utils.ONNX_TO_NUMPY_DTYPE[logit_dtype])).output[0]
shapes = tf_ori_node.output_shapes
dtypes = tf_ori_node.output_dtypes
ctx.remove_node(tf_ori_node.name)
res = ctx.make_node(op_type="Mul", inputs=[log_prob, const_negative_one],
outputs=[tf_ori_node.output[0]], shapes=[shapes[0]], dtypes=[dtypes[0]])
def _make_gathernd_inner_loop(ctx, params, index, dtype):
"""create the inner loop for GatherNd."""
# gather_cur = params
# for (int i = 0; i < size(index); i++)
# gather_res = gather(gather_cur, index[i])
scope_name = utils.make_name("gathernd_inner_loop")
trip_node = ctx.make_node("Size", [index.output[0]])
cond_const = ctx.make_const(utils.make_name("cond"), np.ones((), dtype=np.bool))
trip_name = utils.make_name("i")
cond_name = utils.make_name("cond")
cond_out_name = utils.make_name("cond_out")
cur_name = utils.make_name("gather_cur")
result_name = utils.make_name("res")
# body graph creation
g = ctx.create_new_graph_with_same_config()
g.add_graph_input(trip_name, TensorProto.INT64, [1])
g.add_graph_input(cond_name, TensorProto.BOOL, [])
g.add_graph_input(cur_name, dtype, [])
g.parent_graph = ctx
index_i = g.make_node("Gather", [index.output[0], trip_name], attr={"axis": 0})
false_input_id = gather_input_ids[2]
false_input_id_for_current_iter = get_inputs_for_current_iteration(g, false_input_id, iter_name)
input_ids_for_current_iter = [cond_input_id_for_current_iter, true_input_id_for_current_iter,
false_input_id_for_current_iter]
output_id = None
rank -= 1
if rank >= 1:
loop_1 = create_loop_op(g, input_ids_for_current_iter, output_data_type, output_shape[1:],
trip_count_input_ids, rank)
output_id = loop_1.output[1]
elif rank == 0:
_, if_node_output_id = create_if_op(g, input_ids_for_current_iter, output_data_type, output_shape[1:])
output_id = if_node_output_id
output_identity_name = utils.make_name("loop_output")
loop_output_id = utils.port_name(output_identity_name)
g.make_node(
'Identity',
[output_id],
outputs=[loop_output_id],
name=output_identity_name
)
cond_identity_name = utils.make_name("cond_output")
cond_output_id = utils.port_name(cond_identity_name)
g.make_node(
'Identity',
[cond_name],
outputs=[cond_output_id],
name=cond_identity_name
)
output_dtype = dtype
logger.debug("insert cast for node %s on input %s", op.name, input_name)
if input_node and input_node.type == "Cast" \
and len(g.find_output_consumers(input_node.output[0])) == 1:
input_node.set_attr("to", onnx_pb.TensorProto.FLOAT)
g.set_dtype(input_name, onnx_pb.TensorProto.FLOAT)
else:
cast_node = g.insert_new_node_on_input(op, "Cast", input_name)
cast_node.set_attr("to", onnx_pb.TensorProto.FLOAT)
g.set_dtype(cast_node.output[0], onnx_pb.TensorProto.FLOAT)
g.copy_shape(input_name, cast_node.output[0])
cast_inserted.append(cast_node)
if output_dtype:
# insert reverse cast if needed
for output_name in op.output:
name = utils.make_name(op.name)
logger.debug("insert cast back for node %s on output %s [dtype=%s]", op.name, output_name,
output_dtype)
output_cast = g.insert_new_node_on_output("Cast", output_name, name=name)
output_cast.set_attr("to", output_dtype)
g.set_dtype(output_cast.output[0], output_dtype)
g.copy_shape(output_name, output_cast.output[0])
cast_inserted.append(output_cast)
if cast_inserted:
new_ops.extend(cast_inserted)
new_ops.append(op)
return new_ops