def simple_if(condition_val):
condition = ng.constant(condition_val, dtype=np.bool)
# then_body
X_t = ng.parameter([2], np.float32, "X")
Y_t = ng.parameter([2], np.float32, "Y")
then_mul = ng.multiply(X_t, Y_t)
then_body_res_1 = ng.result(then_mul)
then_body = GraphBody([X_t, Y_t], [then_body_res_1])
then_body_inputs = [
TensorIteratorInvariantInputDesc(1, 0),
TensorIteratorInvariantInputDesc(2, 1)
]
then_body_outputs = [TensorIteratorBodyOutputDesc(0, 0)]
# else_body
X_e = ng.parameter([2], np.float32, "X")
Y_e = ng.parameter([2], np.float32, "Y")
add_e = ng.add(X_e, Y_e)
else_body_res_1 = ng.result(add_e)
else_body = GraphBody([X_e, Y_e], [else_body_res_1])
else_body_inputs = [
TensorIteratorInvariantInputDesc(1, 0),
TensorIteratorInvariantInputDesc(2, 1)
]
else_body_outputs = [TensorIteratorBodyOutputDesc(0, 0)]
X = ng.constant([3, 4], dtype=np.float32)
Y = ng.constant([2, 1], dtype=np.float32)
if_node = ng.if_op(condition, [X, Y], (then_body, else_body),
(then_body_inputs, else_body_inputs),
(then_body_outputs, else_body_outputs))
relu = ng.relu(if_node)
return relu
def create_simple_if_with_two_outputs(condition_val):
condition = ng.constant(condition_val, dtype=np.bool)
# then_body
X_t = ng.parameter([], np.float32, "X")
Y_t = ng.parameter([], np.float32, "Y")
Z_t = ng.parameter([], np.float32, "Z")
add_t = ng.add(X_t, Y_t)
mul_t = ng.multiply(Y_t, Z_t)
then_body_res_1 = ng.result(add_t)
then_body_res_2 = ng.result(mul_t)
then_body = GraphBody([X_t, Y_t, Z_t], [then_body_res_1, then_body_res_2])
then_body_inputs = [
TensorIteratorInvariantInputDesc(1, 0),
TensorIteratorInvariantInputDesc(2, 1),
TensorIteratorInvariantInputDesc(3, 2)
]
then_body_outputs = [
TensorIteratorBodyOutputDesc(0, 0),
TensorIteratorBodyOutputDesc(1, 1)
]
# else_body
X_e = ng.parameter([], np.float32, "X")
Z_e = ng.parameter([], np.float32, "Z")
W_e = ng.parameter([], np.float32, "W")
add_e = ng.add(X_e, W_e)
pow_e = ng.power(W_e, Z_e)
else_body_res_1 = ng.result(add_e)
else_body_res_2 = ng.result(pow_e)
else_body = GraphBody([X_e, Z_e, W_e], [else_body_res_1, else_body_res_2])
else_body_inputs = [
TensorIteratorInvariantInputDesc(1, 0),
TensorIteratorInvariantInputDesc(3, 1),
TensorIteratorInvariantInputDesc(4, 2)
]
else_body_outputs = [
TensorIteratorBodyOutputDesc(0, 0),
TensorIteratorBodyOutputDesc(1, 1)
]
X = ng.constant(15.0, dtype=np.float32)
Y = ng.constant(-5.0, dtype=np.float32)
Z = ng.constant(4.0, dtype=np.float32)
W = ng.constant(2.0, dtype=np.float32)
if_node = ng.if_op(condition, [X, Y, Z, W], (then_body, else_body),
(then_body_inputs, else_body_inputs),
(then_body_outputs, else_body_outputs))
return if_node
def loop(
trip_count: NodeInput,
execution_condition: NodeInput,
inputs: List[Node],
graph_body: GraphBody,
slice_input_desc: List[TensorIteratorSliceInputDesc],
merged_input_desc: List[TensorIteratorMergedInputDesc],
invariant_input_desc: List[TensorIteratorInvariantInputDesc],
body_output_desc: List[TensorIteratorBodyOutputDesc],
concat_output_desc: List[TensorIteratorConcatOutputDesc],
body_condition_output_idx: int,
current_iteration_input_idx: int = -1,
name: Optional[str] = None,
) -> Node:
"""Perform recurrent execution of the network described in the body, iterating through the data.
:param trip_count: A scalar or 1D tensor with 1 element specifying
maximum number of iterations.
:param execution_condition: A scalar or 1D tensor with 1 element
specifying whether to execute the first iteration or not.
:param inputs: The provided to TensorIterator operator.
:param graph_body: The graph representing the body we execute.
:param slice_input_desc: The descriptors describing sliced inputs, that is nodes
representing tensors we iterate through, processing single
data slice in one iteration.
:param merged_input_desc: The descriptors describing merged inputs, that is nodes
representing variables with initial value at first iteration,
which may be changing through iterations.
:param invariant_input_desc: The descriptors describing invariant inputs, that is nodes
representing variable with persistent value through all
iterations.
:param body_output_desc: The descriptors describing body outputs from specified
iteration.
:param concat_output_desc: The descriptors describing specified output values through
all the iterations concatenated into one node.
:param body_condition_output_idx: Determines the purpose of the corresponding result in
the graph_body. This result will determine the dynamic
exit condition. If the value of this result is False,
then iterations stop.
:param current_iteration_input_idx: Determines the purpose of the corresponding parameter
in the graph_body. This parameter will be used as
an iteration counter. Optional.
:return: The new node which performs Loop.
"""
attributes = {
"body": graph_body.serialize(),
"input_descriptions": {"slice_input_desc": [desc.serialize() for desc in slice_input_desc],
"merged_input_desc": [desc.serialize() for desc in merged_input_desc],
"invariant_input_desc": [desc.serialize() for desc in invariant_input_desc]},
"output_descriptions": {"body_output_desc": [desc.serialize() for desc in body_output_desc],
"concat_output_desc": [desc.serialize() for desc in concat_output_desc]},
"special_body_ports": {"body_condition_output_idx": body_condition_output_idx,
"current_iteration_input_idx": current_iteration_input_idx}
}
return _get_node_factory_opset5().create("Loop", as_nodes(trip_count, execution_condition, *inputs),
attributes)
def create_diff_if_with_two_outputs(condition_val):
condition = ng.constant(condition_val, dtype=np.bool)
# then_body
X_t = ng.parameter([2], np.float32, "X")
Y_t = ng.parameter([2], np.float32, "Y")
mmul_t = ng.matmul(X_t, Y_t, False, False)
mul_t = ng.multiply(Y_t, X_t)
then_body_res_1 = ng.result(mmul_t)
then_body_res_2 = ng.result(mul_t)
then_body = GraphBody([X_t, Y_t], [then_body_res_1, then_body_res_2])
then_body_inputs = [
TensorIteratorInvariantInputDesc(1, 0),
TensorIteratorInvariantInputDesc(2, 1)
]
then_body_outputs = [
TensorIteratorBodyOutputDesc(0, 0),
TensorIteratorBodyOutputDesc(1, 1)
]
# else_body
X_e = ng.parameter([2], np.float32, "X")
Z_e = ng.parameter([], np.float32, "Z")
mul_e = ng.multiply(X_e, Z_e)
else_body_res_1 = ng.result(Z_e)
else_body_res_2 = ng.result(mul_e)
else_body = GraphBody([X_e, Z_e], [else_body_res_1, else_body_res_2])
else_body_inputs = [
TensorIteratorInvariantInputDesc(1, 0),
TensorIteratorInvariantInputDesc(3, 1)
]
else_body_outputs = [
TensorIteratorBodyOutputDesc(0, 0),
TensorIteratorBodyOutputDesc(1, 1)
]
X = ng.constant([3, 4], dtype=np.float32)
Y = ng.constant([2, 1], dtype=np.float32)
Z = ng.constant(4.0, dtype=np.float32)
if_node = ng.if_op(condition, [X, Y, Z], (then_body, else_body),
(then_body_inputs, else_body_inputs),
(then_body_outputs, else_body_outputs))
return if_node
def simple_if_without_parameters(condition_val):
condition = ng.constant(condition_val, dtype=np.bool)
# then_body
then_constant = ng.constant(0.7, dtype=np.float)
then_body_res_1 = ng.result(then_constant)
then_body = GraphBody([], [then_body_res_1])
then_body_inputs = []
then_body_outputs = [TensorIteratorBodyOutputDesc(0, 0)]
# else_body
else_const = ng.constant(9.0, dtype=np.float)
else_body_res_1 = ng.result(else_const)
else_body = GraphBody([], [else_body_res_1])
else_body_inputs = []
else_body_outputs = [TensorIteratorBodyOutputDesc(0, 0)]
if_node = ng.if_op(condition, [], (then_body, else_body),
(then_body_inputs, else_body_inputs),
(then_body_outputs, else_body_outputs))
relu = ng.relu(if_node)
return relu
def test_tensor_iterator():
from ngraph.utils.tensor_iterator_types import (
GraphBody,
TensorIteratorSliceInputDesc,
TensorIteratorMergedInputDesc,
TensorIteratorInvariantInputDesc,
TensorIteratorBodyOutputDesc,
TensorIteratorConcatOutputDesc,
)
# Body parameters
body_timestep = ng.parameter([], np.int32, "timestep")
body_data_in = ng.parameter([1, 2, 2], np.float32, "body_in")
body_prev_cma = ng.parameter([2, 2], np.float32, "body_prev_cma")
body_const_one = ng.parameter([], np.int32, "body_const_one")
# CMA = cumulative moving average
prev_cum_sum = ng.multiply(ng.convert(body_timestep, "f32"), body_prev_cma)
curr_cum_sum = ng.add(prev_cum_sum, ng.squeeze(body_data_in, [0]))
elem_cnt = ng.add(body_const_one, body_timestep)
curr_cma = ng.divide(curr_cum_sum, ng.convert(elem_cnt, "f32"))
cma_hist = ng.unsqueeze(curr_cma, [0])
# TI inputs
data = ng.parameter([16, 2, 2], np.float32, "data")
# Iterations count
zero = ng.constant(0, dtype=np.int32)
one = ng.constant(1, dtype=np.int32)
initial_cma = ng.constant(np.zeros([2, 2], dtype=np.float32),
dtype=np.float32)
iter_cnt = ng.range(zero, np.int32(16), np.int32(1))
ti_inputs = [iter_cnt, data, initial_cma, one]
graph_body = GraphBody(
[body_timestep, body_data_in, body_prev_cma, body_const_one],
[curr_cma, cma_hist])
ti_slice_input_desc = [
# timestep
# input_idx, body_param_idx, start, stride, part_size, end, axis
TensorIteratorSliceInputDesc(0, 0, 0, 1, 1, -1, 0),
# data
TensorIteratorSliceInputDesc(1, 1, 0, 1, 1, -1, 0),
]
ti_merged_input_desc = [
# body prev/curr_cma
TensorIteratorMergedInputDesc(2, 2, 0),
]
ti_invariant_input_desc = [
# body const one
TensorIteratorInvariantInputDesc(3, 3),
]
# TI outputs
ti_body_output_desc = [
# final average
TensorIteratorBodyOutputDesc(0, 0, -1),
]
ti_concat_output_desc = [
# history of cma
TensorIteratorConcatOutputDesc(1, 1, 0, 1, 1, -1, 0),
]
node = ng.tensor_iterator(
ti_inputs,
graph_body,
ti_slice_input_desc,
ti_merged_input_desc,
ti_invariant_input_desc,
ti_body_output_desc,
ti_concat_output_desc,
)
assert node.get_type_name() == "TensorIterator"
assert node.get_output_size() == 2
# final average
assert list(node.get_output_shape(0)) == [2, 2]
# cma history
assert list(node.get_output_shape(1)) == [16, 2, 2]
