def test_unsqueeze():
runtime = get_runtime()
data_shape = [3, 4, 5]
parameter_data = ov.parameter(data_shape, name="Data", dtype=np.float32)
data_value = np.arange(60.0, dtype=np.float32).reshape(3, 4, 5)
axes = [0, 4]
model = ov.unsqueeze(parameter_data, axes)
computation = runtime.computation(model, parameter_data)
result = computation(data_value)
expected = np.arange(60.0, dtype=np.float32).reshape([1, 3, 4, 5, 1])
assert np.allclose(result, expected)
def test_loop_basic():
bool_val = np.array([1], dtype=np.bool)
bool_val[0] = True
condition = ov.constant(bool_val)
trip_count = ov.constant(16, dtype=np.int32)
# Body parameters
body_timestep = ov.parameter([], np.int32, "timestep")
body_data_in = ov.parameter([1, 2, 2], np.float32, "body_in")
body_prev_cma = ov.parameter([2, 2], np.float32, "body_prev_cma")
body_const_one = ov.parameter([], np.int32, "body_const_one")
body_ports = [-1, 2]
# CMA = cumulative moving average
prev_cum_sum = ov.multiply(ov.convert(body_timestep, "f32"), body_prev_cma)
curr_cum_sum = ov.add(prev_cum_sum, ov.squeeze(body_data_in, [0]))
elem_cnt = ov.add(body_const_one, body_timestep)
curr_cma = ov.divide(curr_cum_sum, ov.convert(elem_cnt, "f32"))
cma_hist = ov.unsqueeze(curr_cma, [0])
# TI inputs
data = ov.parameter([16, 2, 2], np.float32, "data")
# Iterations count
zero = ov.constant(0, dtype=np.int32)
one = ov.constant(1, dtype=np.int32)
initial_cma = ov.constant(np.zeros([2, 2], dtype=np.float32),
dtype=np.float32)
iter_cnt = ov.range(zero, np.int32(16), np.int32(1))
body_const_condition = ov.constant(bool_val)
graph_body = Model(
[curr_cma, cma_hist, body_const_condition],
[body_timestep, body_data_in, body_prev_cma, body_const_one],
"body_function")
ti_slice_input_desc = [
# timestep
# input_idx, body_param_idx, start, stride, part_size, end, axis
SliceInputDescription(2, 0, 0, 1, 1, -1, 0),
# data
SliceInputDescription(3, 1, 0, 1, 1, -1, 0),
]
ti_merged_input_desc = [
# body prev/curr_cma
MergedInputDescription(4, 2, 0),
]
ti_invariant_input_desc = [
# body const one
InvariantInputDescription(5, 3),
]
# TI outputs
ti_body_output_desc = [
# final average
BodyOutputDescription(0, 0, -1),
]
ti_concat_output_desc = [
# history of cma
ConcatOutputDescription(1, 1, 0, 1, 1, -1, 0),
]
loop = ov.loop(trip_count, condition)
loop.set_function(graph_body)
loop.set_special_body_ports(body_ports)
loop.set_sliced_input(body_timestep, iter_cnt.output(0), 0, 1, 1, 0, 0)
loop.set_sliced_input(body_data_in, data.output(0), 0, 1, 1, 0, 0)
# set different end parameter
loop.set_input_descriptions(ti_slice_input_desc)
loop.set_merged_input(body_prev_cma, initial_cma.output(0),
curr_cma.output(0))
loop.set_invariant_input(body_const_one, one.output(0))
loop.get_iter_value(curr_cma.output(0), -1)
loop.get_concatenated_slices(cma_hist.output(0), 0, 1, 1, -1, 0)
assert loop.get_function() == graph_body
assert loop.get_special_body_ports() == body_ports
assert loop.get_num_iterations() == 16
input_desc = loop.get_input_descriptions()
output_desc = loop.get_output_descriptions()
assert len(input_desc) == len(ti_slice_input_desc) + \
len(ti_merged_input_desc) + len(ti_invariant_input_desc)
assert len(
output_desc) == len(ti_body_output_desc) + len(ti_concat_output_desc)
for i in range(len(ti_slice_input_desc)):
assert input_desc[i].get_type_info(
) == ti_slice_input_desc[i].get_type_info()
assert input_desc[i].input_index == ti_slice_input_desc[i].input_index
assert input_desc[i].body_parameter_index == ti_slice_input_desc[
i].body_parameter_index
assert input_desc[i].start == ti_slice_input_desc[i].start
assert input_desc[i].stride == ti_slice_input_desc[i].stride
assert input_desc[i].part_size == ti_slice_input_desc[i].part_size
assert input_desc[i].end == ti_slice_input_desc[i].end
assert input_desc[i].axis == ti_slice_input_desc[i].axis
for i in range(len(ti_merged_input_desc)):
assert input_desc[
len(ti_slice_input_desc) +
i].get_type_info() == ti_merged_input_desc[i].get_type_info()
assert input_desc[len(ti_slice_input_desc) +
i].input_index == ti_merged_input_desc[i].input_index
assert input_desc[len(ti_slice_input_desc) +
i].body_parameter_index == ti_merged_input_desc[
i].body_parameter_index
assert input_desc[
len(ti_slice_input_desc) +
i].body_value_index == ti_merged_input_desc[i].body_value_index
for i in range(len(ti_concat_output_desc)):
assert output_desc[
len(ti_body_output_desc) +
i].get_type_info() == ti_concat_output_desc[i].get_type_info()
assert output_desc[
len(ti_body_output_desc) +
i].output_index == ti_concat_output_desc[i].output_index
assert output_desc[
len(ti_body_output_desc) +
i].body_value_index == ti_concat_output_desc[i].body_value_index
assert output_desc[len(ti_body_output_desc) +
i].start == ti_concat_output_desc[i].start
assert output_desc[len(ti_body_output_desc) +
i].stride == ti_concat_output_desc[i].stride
assert output_desc[len(ti_body_output_desc) +
i].part_size == ti_concat_output_desc[i].part_size
assert output_desc[len(ti_body_output_desc) +
i].end == ti_concat_output_desc[i].end
assert output_desc[len(ti_body_output_desc) +
i].axis == ti_concat_output_desc[i].axis