def test_single_input_to_single_operation(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
def prog(x):
x = mb.square(x=x)
return x
self.assertEqual(get_op_types_in_program(prog), ['square'])
apply_pass_and_basic_check(
prog, transform.FP16ComputePrecision(op_selector=lambda op: True))
_, _, block = apply_pass_and_basic_check(
prog, "common::dead_code_elimination")
self.assertEqual(get_op_types_in_program(prog),
["cast", "square", "cast"])
# Asserting first cast configuration
cast_1 = block.find_ops(op_type="cast")[0]
self.assertEqual(cast_1.dtype.val, "fp16")
self.assertEqual(len(cast_1.outputs), 1)
self.assertEqual(len(cast_1.outputs[0].child_ops), 1)
self.assertEqual(cast_1.outputs[0].child_ops[0].op_type, "square")
# Asserting second cast configuration
cast_2 = block.find_ops(op_type="cast")[1]
self.assertEqual(cast_2.dtype.val, "fp32")
self.assertEqual(len(cast_2.outputs), 1)
self.assertEqual(len(cast_2.outputs[0].child_ops), 0)
assert_model_is_valid(
prog,
{"x": (10, 20)},
expected_output_shapes={block.outputs[0].name: (10, 20)},
)
def test_divide_by_zero_operation(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
def prog(x):
eps = mb.const(val=1e-10)
x = mb.real_div(x=x, y=eps)
return x
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, transform.FP16ComputePrecision(op_selector=lambda op: True))
mlmodel = ct.convert(prog, source="milinternal")
input_dict = {"x": np.random.rand(10, 20)}
if _IS_MACOS:
prediction = mlmodel.predict(input_dict, useCPUOnly=True)
assert (not np.isnan(prediction['real_div_0']).any())
assert (np.isfinite(prediction['real_div_0']).all())