def _add_batch_norm(x, mean, variance, scale, offset, epsilon, name):
if mean.shape[0] != 0 and variance.shape[0] != 0:
# In this case, we can use the mb.batch_norm directly
x = mb.batch_norm(x=x,
mean=mean,
variance=variance,
gamma=scale,
beta=offset,
epsilon=epsilon,
name=name)
else:
# In this case, we need to manually compute the batch_norm
axes = [axis for axis in range(x.rank) if axis != 1]
mean = mb.reduce_mean(x=x, axes=axes, keep_dims=True)
num = mb.sub(x=x, y=mean)
square = mb.mul(x=num, y=num)
variance = mb.reduce_mean(x=square, axes=axes, keep_dims=True)
variance_add_epsilon = mb.add(x=variance, y=epsilon)
sqrt = mb.sqrt(x=variance_add_epsilon)
x = mb.real_div(x=num, y=sqrt)
shape = [1] * x.rank
shape[1] = -1 if any_symbolic(scale.shape) else scale.shape[0]
scale_reshape = mb.reshape(x=scale, shape=shape)
offset_reshape = mb.reshape(x=offset, shape=shape)
x = mb.mul(x=x, y=scale_reshape)
x = mb.add(x=x, y=offset_reshape, name=name)
return x
def ops_arrangement(x):
power = mb.pow(x=x, y=3, name="thepowerop")
sub_0 = mb.sub(x=power, y=5, name="sub_0")
add_0 = mb.add(x=power, y=5, name="add_0")
mul_0 = mb.mul(x=power, y=5, name="mul_0")
add_1 = mb.add(x=add_0, y=mul_0, name="add_1")
add_2 = mb.add(x=sub_0, y=add_1,name="add_2")
return add_2
def prog(x):
x1 = mb.pow(x=x, y=3)
x1 = mb.mul(x=0.044715, y=x1)
x1 = mb.add(x=x1, y=x)
x1 = mb.mul(x=x1, y=np.sqrt(2 / np.pi))
x1 = mb.tanh(x=x1)
x1 = mb.add(x=1, y=x1)
x1 = mb.mul(x=0.5, y=x1)
x1 = mb.mul(x=x, y=x1)
return x1
def prog(x):
firstmul = mb.mul(x=x, y=0.5) if first_op_1 else mb.mul(x=0.5, y=x)
x1 = mb.pow(x=x, y=3)
x1 = mb.mul(x=0.044715, y=x1) if first_op_2 else mb.mul(x=x1, y=0.044715)
x1 = mb.add(x=x1, y=x) if first_op_3 else mb.add(x=x, y=x1)
x1 = mb.mul(x=x1, y=np.sqrt(2 / np.pi)) if first_op_4 else mb.mul(x=np.sqrt(2 / np.pi), y=x1)
x1 = mb.tanh(x=x1)
x1 = mb.add(x=1, y=x1) if first_op_5 else mb.add(x=x1, y=1)
x1 = mb.mul(x=firstmul, y=x1) if first_op_6 else mb.mul(x=x1, y=firstmul)
return x1
def _insert_image_preprocessing_ops(block, prog):
input_types = list(prog.main_input_types)
for input_type in input_types:
if isinstance(input_type, ImageType):
if input_type.name not in block.inputs:
continue
input_var = block.inputs[input_type.name]
placeholder_op = block.placeholder_inputs[input_type.name]
first_op = block.operations[0]
old_var = placeholder_op.outputs[0]
has_bias = np.any(np.array(input_type.bias) != 0)
with block:
last_output = input_var
input_nptype = nptype_from_builtin(type(last_output.dtype()))
if input_type.scale != 1:
last_output = mb.mul(x=last_output,
y=np.array(input_type.scale,
dtype=input_nptype),
before_op=first_op,
name=input_var.name + "__scaled__")
if has_bias:
if input_type.color_layout == "G":
last_output = mb.add(x=last_output,
y=np.array(input_type.bias,
dtype=input_nptype),
before_op=first_op,
name=input_var.name +
"__biased__")
else:
if len(last_output.shape) == 3:
last_output = mb.add(
x=last_output,
y=np.array(input_type.bias,
dtype=input_nptype).reshape(
[3, 1, 1]),
before_op=first_op,
name=input_var.name + "__biased__")
elif len(last_output.shape) == 4:
last_output = mb.add(
x=last_output,
y=np.array(input_type.bias,
dtype=input_nptype).reshape(
[1, 3, 1, 1]),
before_op=first_op,
name=input_var.name + "__biased__")
else:
raise TypeError(
"Unsupported rank for image input type.")
if last_output != input_var:
block.replace_uses_of_var_after_op(anchor_op=last_output.op,
old_var=old_var,
new_var=last_output)
def prog(x):
x = mb.transpose(x=x, perm=[0, 3, 1, 2])
if flip_mul_input_order:
x = mb.mul(x=gamma, y=x)
else:
x = mb.mul(x=x, y=gamma)
if flip_add_input_order:
x = mb.add(x=beta, y=x)
else:
x = mb.add(x=x, y=beta)
return x
def prog(x):
mean0 = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=True)
sub = mb.sub(x=x, y=mean0)
square = mb.square(x=sub)
mean1 = mb.reduce_mean(x=square, axes=[2, 3], keep_dims=True)
add_eps = mb.add(x=mean1, y=1e-5) # epsilon
rsqrt = mb.rsqrt(x=add_eps)
mul1 = mb.mul(x=rsqrt, y=x)
mul2 = mb.mul(x=mean0, y=rsqrt)
mul_sub = mb.mul(x=mul2, y=-1)
add = mb.add(x=mul1, y=mul_sub)
return add
def prog(x):
mean0 = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=True)
sub0 = mb.sub(x=x, y=mean0)
sub1 = mb.sub(x=x, y=mean0)
square = mb.square(x=sub0)
mean1 = mb.reduce_mean(x=square, axes=[2, 3], keep_dims=True)
add_eps = mb.add(x=mean1, y=1e-5) # epsilon
pow = mb.pow(x=add_eps, y=0.5)
div = mb.real_div(x=sub1, y=pow)
mul_gamma = mb.mul(x=np.random.rand(1, shape[1], 1, 1), y=div) #
add_beta = mb.add(x=np.random.rand(1, shape[1], 1, 1), y=mul_gamma)
return add_beta
def prog(x):
x1 = mb.reduce_mean(x=x, axes=axes, keep_dims=True)
x2 = mb.sub(x=x, y=x1)
x2 = mb.square(x=x2)
x2 = mb.reduce_mean(x=x2, axes=axes, keep_dims=True)
x2 = mb.add(x=x2, y=1e-5)
x2 = mb.rsqrt(x=x2)
x3 = mb.mul(x=np.random.rand(*shape[-len(axes):]), y=x2)
x4 = mb.mul(x=x3, y=x1)
x5 = mb.mul(x=x, y=x3)
x4 = mb.sub(x=np.random.rand(*shape[-len(axes):]), y=x4)
y = mb.add(x=x4, y=x5)
return y
def prog(x):
x1 = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=True)
x2 = mb.sub(x=x, y=x1)
x2 = mb.square(x=x2)
x2 = mb.reduce_mean(x=x2, axes=[2, 3], keep_dims=True)
x2 = mb.add(x=x2, y=1e-5)
x2 = mb.rsqrt(x=x2)
x3 = mb.mul(x=np.random.rand(1, shape[1], 1, 1), y=x2)
x4 = mb.mul(x=x3, y=x1)
x5 = mb.mul(x=x, y=x3)
x4 = mb.sub(x=np.random.rand(1, shape[1], 1, 1), y=x4)
y = mb.add(x=x4, y=x5)
return y
def prog(x):
x = mb.transpose(x=x, perm=[0, 2, 3, 1])
x1 = mb.reduce_mean(x=x, axes=[1, 2], keep_dims=True)
x2 = mb.sub(x=x, y=x1)
x2 = mb.square(x=x2)
x2 = mb.reduce_mean(x=x2, axes=[1, 2], keep_dims=True)
x2 = mb.add(x=x2, y=1e-5)
x2 = mb.rsqrt(x=x2)
x3 = mb.mul(x=np.random.rand(1, 1, 1, shape[1]), y=x2)
x4 = mb.mul(x=x3, y=x1)
x5 = mb.mul(x=x, y=x3)
x4 = mb.sub(x=np.random.rand(1, 1, 1, shape[1]), y=x4)
x6 = mb.add(x=x4, y=x5)
y = mb.transpose(x=x6, perm=[0, 3, 1, 2])
return y
def prog(x1, x2):
a = mb.stack(values=[x1, x2], axis=1)
b = mb.reshape(x=a, shape=[-1, 10, 4, 3])
c = mb.relu(x=a)
c = mb.reshape(x=c, shape=[-1, 10, 4, 3])
d = mb.add(x=b, y=c)
return d
def prog(x):
mul_square1 = mb.mul(x=x, y=x)
sum = mb.reduce_sum(x=x, axes=[2, 3], keep_dims=True)
mul_mean = mb.mul(x=sum, y=3.3333334e-05) # dummy value here
mul_square = mb.mul(x=mul_mean, y=mul_mean)
sum1 = mb.reduce_sum(x=mul_square1, axes=[2, 3], keep_dims=True)
mul_mean1 = mb.mul(x=sum1, y=8.333333e-06) # dummy value here
sub_variance = mb.sub(x=mul_mean1, y=mul_square)
add_eps = mb.add(x=sub_variance, y=1e-5) # epsilon
rsqrt = mb.rsqrt(x=add_eps)
mul_gamma = mb.mul(x=rsqrt, y=np.random.rand(1, shape[1], 1, 1))
mul1 = mb.mul(x=mul_gamma, y=x)
mul2 = mb.mul(x=mul_mean, y=mul_gamma)
sub_beta = mb.sub(x=np.random.rand(1, shape[1], 1, 1), y=mul2)
add = mb.add(x=mul1, y=sub_beta)
return add
def prog(x):
s1, s2 = mb.split(x=x, num_splits=2, axis=0)
x2 = mb.square(x=x)
x3 = mb.relu(x=x2)
s1_1 = mb.square(x=s1)
s3 = mb.add(x=s1_1, y=s2)
return x3, s3
def build(x):
lut_data = np.array([
-19.0,
4.0,
0.0,
-1.0,
1.0,
3.0,
5.0,
-8.0,
19,
13,
42,
4.5,
5.4,
2.0,
-6,
-7,
]).astype(np.float32)
indices = np.array([212, 21]).astype(np.uint8)
shape = np.array([4, 1]).astype(np.uint32)
y = mb.constexpr_lut_to_dense(lut=lut_data,
indices=indices,
shape=shape)
return mb.add(x=x, y=y)
def test_nn_backend_style_sanitization(self):
'''
Test that intermediate var names are unchanged, and
only model input and output names are modified, i.e.
sanitized (adhering to the format [a-zA-Z_][a-zA-Z0-9_]*)
for the NN backend.
'''
prog = Program()
func_inputs = {"x/0": mb.placeholder(shape=[2, 3]),
"y": mb.placeholder(shape=[2, 3])}
with Function(func_inputs) as ssa_fun:
x, y = ssa_fun.inputs["x/0"], ssa_fun.inputs["y"]
x = mb.relu(x=x, name="relu/1")
z = mb.add(x=x, y=y, name="out/1")
ssa_fun.set_outputs([z])
prog.add_function("main", ssa_fun)
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::sanitize_input_output_names",
skip_output_name_check=True
)
relu_op = prog.find_ops(op_type="relu", exactly_one=True)[0]
assert relu_op.inputs["x"].name == "x_0" # input name: sanitized
assert relu_op.outputs[0].name == "relu/1" # intermediate name: unchanged
assert block.outputs[0].name == "out_1" # output name: sanitized
# convert prev_prog to NN backend
mlmodel = ct.convert(prev_prog)
spec = mlmodel._spec
assert spec.description.input[0].name == "x_0"
assert spec.description.output[0].name == "out_1"
relu_layer = spec.neuralNetwork.layers[0]
assert relu_layer.output[0] == "relu/1"
def body(i, num_iters, ls, update):
new_elem = mb.mul(x=update, y=i)
return (
mb.add(x=i, y=1),
num_iters,
mb.list_write(ls=ls, index=i, value=new_elem),
update,
)
def prog(x):
weights_val = np.random.rand(2, 4).T.astype(np.float32)
weights = mb.const(val=weights_val, mode="immediate_value")
bias_val = np.random.rand(2).astype(np.float32)
bias = mb.const(val=bias_val, mode="immediate_value")
matmul = mb.matmul(x=x, y=weights)
return mb.add(x=matmul, y=bias)
def build(x):
nonzero_data = np.array([1, 2, 4]).astype(np.float32)
mask = np.array([11]).astype(np.uint8)
shape = np.array([4, 1]).astype(np.uint32)
y = mb.constexpr_sparse_to_dense(nonzero_data=nonzero_data,
mask=mask,
shape=shape)
return mb.add(x=x, y=y)
def build(x):
return [
mb.add(x=x, y=x),
mb.random_bernoulli(shape=np.array([2, 1, 3], np.int32),
prob=1.0),
mb.random_bernoulli(shape=np.array([3, 1, 2], np.int32),
prob=0.0),
]
def test_builder_add(self):
x = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
y = np.array([[-1, 2, -3], [4, -5, 6]], dtype=np.float32)
expected_outputs = np.array([[0, 4, 0], [8, 0, 12]], dtype=np.float32)
v = mb.add(x=x, y=y)
np.testing.assert_allclose(expected_outputs,
v.val,
atol=1e-04,
rtol=1e-05)
def build(x):
return [
mb.add(x=x, y=x),
mb.random_uniform(
shape=np.array([2, 1, 3], np.int32), low=0.0, high=0.0
),
mb.random_uniform(
shape=np.array([3, 1, 2], np.int32), low=1.0, high=1.0
),
]
def build(x):
return [
mb.add(x=x, y=x),
mb.random_normal(
shape=np.array([2, 1, 3], np.int32), mean=1.0, stddev=0.0
),
mb.random_normal(
shape=np.array([3, 1, 2], np.int32), mean=0.0, stddev=0.0
),
]
def prog(x):
if op_type == "real_div":
x1 = mb.real_div(x=x, y=2**0.5)
elif op_type == "mul":
x1 = mb.mul(x=x, y=2**-0.5) if is_first_op1 else mb.mul(x=2**-0.5, y=x)
x2 = mb.erf(x=x1)
x3 = mb.add(x=x2, y=1) if is_first_op2 else mb.add(x=1, y=x2)
if const_mul_first:
y1 = mb.const(val=0.5)
y2 = x
else:
y1 = x
y2 = mb.const(val=0.5)
x4 = mb.mul(x=x3, y=y1) if is_first_op3 else mb.mul(x=y1, y=x3)
x5 = mb.mul(x=x4, y=y2) if is_first_op4 else mb.mul(x=y2, y=x4)
return x5
def program(x):
x = mb.slice_by_index(x=x, begin=[0], end=[1], squeeze_mask=[True])
func = getattr(mb, elem_op)
if reverse_order:
x = func(x=2.0, y=x)
else:
x = func(x=x, y=2.0)
expand = mb.expand_dims(x=x, axes=[0])
other_1 = mb.add(x=x, y=[1, 2, 3])
other_2 = mb.sub(x=x, y=[1, 2, 3])
return expand, other_1, other_2
def build(x):
quantized_data = np.array([3, 5, 5, 6]).reshape(1, 1, 2,
2).astype(np.uint8)
scale = np.array([1, 2]).astype(np.float32)
zero_point = np.array([2, 4]).astype(np.uint8)
axis = 3
y = mb.constexpr_affine_dequantize(
quantized_data=quantized_data,
zero_point=zero_point,
scale=scale,
axis=axis,
)
return mb.add(x=x, y=y)
def body(i, cs_list, h_list):
xi = mb.gather(x=x, indices=i, axis=0)
h_prev = mb.gather(x=h_list, indices=i, axis=0)
cs_prev = mb.gather(x=cs_list, indices=i, axis=0)
ig, cs, fg, og, ci, co, h = _lstm_cell_builder(op, xi, h_prev, cs_prev)
counter = mb.add(x=i, y=1)
return (
counter,
mb.scatter(data=cs_list, indices=counter, updates=cs),
mb.scatter(data=h_list, indices=counter, updates=h),
)
def prog(x):
x = mb.cast(x=x, dtype="fp16")
x1 = mb.square(x=x)
x1_t = mb.transpose(x=x1, perm=[1, 0])
def true_fn():
return mb.add(x=x1_t, y=1, name='x2')
def false_fn():
return mb.add(x=x1_t, y=2, name='x2')
is_one = mb.equal(x=mb.squeeze(x=x), y=1)
pred = mb.squeeze(x=is_one)
x3 = mb.cond(pred=pred, _true_fn=true_fn, _false_fn=false_fn)
x4 = mb.add(x=x1_t, y=x3)
x5 = mb.cast(x=x4, dtype="fp32")
return x5
def prog(x):
kwargs = {
"x": x,
"weight": W,
"pad_type": "valid",
"dilations": [1] * conv_dim,
"strides": [1] * conv_dim,
}
if prebuilt_bias:
kwargs["bias"] = np.random.rand(Cout)
x = mb.conv_transpose(**kwargs) if use_conv_transpose else mb.conv(
**kwargs)
if use_sub_instead:
x = mb.sub(x=x, y=const)
else:
x = mb.add(
x=const if flip_add_input_order else x,
y=x if flip_add_input_order else const,
)
x = mb.relu(x=x)
return x
def body1(i, j):
new_i = mb.while_loop(_cond=cond2,
_body=body2,
loop_vars=(i, ))
return mb.add(x=new_i, y=two), j
