def setup_graph_residual():
"""
generate residual structure
v0 --[op1]--> v1 -+----------------+--[op3]--> v3
| |
+--[op2]--> v2 --+
"""
global graph, op1, op2, op3
global v0, v1, v2, v3
v0 = Variable((1, 1), OrderNC)
op1 = Operator("op1")
v1 = Variable((1, 2), OrderNC)
op2 = TestOperator("op2")
v2 = Variable((1, 3), OrderNC)
op3 = Operator("op3")
v3 = Variable((1, 4), OrderNC)
op1.append_input("v0", v0)
op1.append_output("v1", v1)
op2.append_input("v1", v1)
op2.append_output("v2", v2)
op3.append_input("v1", v1)
op3.append_input("v2", v2)
op3.append_output("v3", v3)
graph = Graph([v0], [v3])
def test_expand_dims_without_axis():
v1 = Variable([2, 1, 1, 3], OrderNHWC)
v2 = v1.squeeze()
assert v2.order == Order([Axis.N, Axis.C])
assert v2.shape_dict[Axis.N] == 2
assert v2.shape_dict[Axis.C] == 3
assert isinstance(v2.output_from, Reshape)
assert v2.output_from.inputs["x"] == v1
def test_transpose_like():
v1 = Variable([2, 3, 4, 5], OrderNHWC)
v2 = Variable([2, 5, 3, 4], OrderNCHW)
v3 = v1.transpose_like(v2)
assert v3.shape == (2, 5, 3, 4), v3.shape
assert v3.order == OrderNCHW
assert isinstance(v3.output_from, Transpose)
assert v3.output_from.inputs["x0"] == v1
def get(base: Variable):
if not base.has_attribute(TextureShape):
attribute = TextureShape(base)
base.attributes.add(attribute)
else:
attribute = base.get_attribute(TextureShape)[0]
return [attribute.height, attribute.width]
def test_reshape():
v1 = Variable([2, 3, 4, 5], OrderNHWC)
v2 = v1.reshape(shape=[1, 6, 4, 5], order=OrderNCHW)
assert v2.shape_dict[Axis.N] == 1
assert v2.shape_dict[Axis.C] == 6
assert v2.shape_dict[Axis.H] == 4
assert v2.shape_dict[Axis.W] == 5
assert isinstance(v2.output_from, Reshape)
assert v2.output_from.inputs["x"] == v1
def get(variable: Variable):
if variable.has_attribute(TextureShape):
attribute = variable.get_attribute(TextureShape)[0]
else:
attribute = TextureShape(variable)
variable.attributes.add(attribute)
return attribute.height, attribute.width
def set(variable: Variable, width: int, height: int):
if variable.has_attribute(TextureShape):
attribute = variable.get_attribute(TextureShape)[0]
else:
attribute = TextureShape(variable)
variable.attributes.add(attribute)
attribute.width = width
attribute.height = height
def test_expand_dims_without_index():
v1 = Variable([2, 3], OrderNC)
v2 = v1.expand_dims(Axis.H)
assert v2.order == Order([Axis.N, Axis.C, Axis.H])
assert v2.shape_dict[Axis.N] == 2
assert v2.shape_dict[Axis.H] == 1
assert v2.shape_dict[Axis.C] == 3
assert isinstance(v2.output_from, Reshape)
assert v2.output_from.inputs["x"] == v1
def test_squeeze_with_one_axis():
v1 = Variable([2, 1, 1, 3], OrderNHWC)
v2 = v1.squeeze(Axis.H)
assert v2.order == Order([Axis.N, Axis.W, Axis.C])
assert v2.shape_dict[Axis.N] == 2
assert v2.shape_dict[Axis.W] == 1
assert v2.shape_dict[Axis.C] == 3
assert isinstance(v2.output_from, Reshape)
assert v2.output_from.inputs["x"] == v1
def set(base: Variable, width: int, height: int):
if not base.has_attribute(TextureShape):
attribute = TextureShape(base)
base.attributes.add(attribute)
else:
attribute = base.get_attribute(TextureShape)[0]
attribute.width = width
attribute.height = height
def main(k, s, p, n, h1, w1, c1, c2, expected_shape_dict: AxisKeyDict[int]):
op = Convolution2D(None, ksize=k, stride=s, padding=p)
x = Variable((n, h1, w1, c1), Order([Axis.N, Axis.H, Axis.W, Axis.C]))
w = Variable((c1, op.ksize[0], op.ksize[1], c2), Order([Axis.C, Axis.KH, Axis.KW, Axis.N]))
y, = op(x, w)
for axis in y.order.axes:
assert y.shape_dict[axis] == expected_shape_dict[axis]
def test_get_input_name():
op = Operator("op")
v1 = Variable((1, 2, 3, 4), OrderNHWC)
v2 = Variable((1, 2, 3, 4), OrderNHWC)
op.append_input("v1", v1)
op.append_input("v2", v2)
assert op.get_input_name(v1) == "v1"
assert op.get_input_name(v2) == "v2"
def test_sub_with_variable():
v1 = Variable([2, 3, 4, 5], OrderNHWC)
v2 = Variable([2, 3, 4, 5], OrderNHWC)
v3 = v1 - v2
assert isinstance(v3.output_from, ElementwiseAdd)
assert v3.output_from.inputs["x0"] == v1
neg_v2 = v3.output_from.inputs["x1"]
assert isinstance(neg_v2.output_from, ScalarMul)
assert neg_v2.output_from.inputs["x0"] == v2
assert neg_v2.output_from.value == -1
def main(k, s, p, n, h1, w1, c1, expected_shape_dict: Dict[Axis, int]):
for order_x in orders4:
op = Col2Im(None, ksize=k, stride=s, padding=p)
x = Variable((n, h1, w1, c1), OrderNHWC)
x.change_order(order_x)
y, = op(x)
for axis in y.order.axes:
assert y.shape_dict[axis] == expected_shape_dict[axis]
def test_combine_axes():
v1 = Variable([2, 3, 4, 5], OrderNHWC)
v2 = v1.combine_axes([Axis.W, Axis.H], Axis.H)
assert v2.order == Order([Axis.N, Axis.H, Axis.C])
assert v2.shape_dict[Axis.N] == 2
assert v2.shape_dict[Axis.H] == 12
assert v2.shape_dict[Axis.C] == 5
assert isinstance(v2.output_from, Reshape)
assert v2.output_from.in_order == Order([Axis.N, Axis.W, Axis.H, Axis.C])
assert v2.output_from.out_order == Order([Axis.N, Axis.H, Axis.C])
assert v2.output_from.inputs["x"] == v1
def test_replace_input():
op = Operator("op")
v1 = Variable((1, 2, 3, 4), OrderNHWC)
v2 = Variable((1, 2, 3, 4), OrderNHWC)
op.append_input("v1", v1)
op.replace_input(v1, v2)
assert op.inputs["v1"] == v2
assert v1.input_to == set()
assert v2.input_to == {op}
def test_replace_output():
op = Operator("op")
v1 = Variable((1, 2, 3, 4), OrderNHWC)
v2 = Variable((1, 2, 3, 4), OrderNHWC)
op.append_output("v1", v1)
op.replace_output(v1, v2)
assert op.outputs["v1"] == v2
assert v1.output_from is None
assert v2.output_from == op
def main(k, s, p, d, n, h1, w1, c1, expected_shape_dict: AxisKeyDict[int]):
for order_x in orders4:
op = Im2Col("im2col", ksize=k, stride=s, padding=p, dilation_rate=d)
x = Variable((n, h1, w1, c1), OrderNHWC)
x.change_order(order_x)
y, = op(x)
for axis in y.order.axes:
assert y.shape_dict[axis] == expected_shape_dict[axis]
def test_append_output():
op = Operator("op")
v1 = Variable((1, 2, 3, 4), OrderNHWC)
v2 = Variable((1, 2, 3, 4), OrderNHWC)
op.append_output("v1", v1)
op.append_output("v2", v2)
assert op.outputs["v1"] == v1
assert op.outputs["v2"] == v2
assert v1.output_from == op
assert v2.output_from == op
def replace_variable(graph: Graph, old_var: Variable, new_var: Variable):
old_var.replace(new_var)
if old_var in graph.inputs:
i = graph.inputs.index(old_var)
graph.inputs.remove(old_var)
graph.inputs.insert(i, new_var)
if old_var in graph.outputs:
i = graph.outputs.index(old_var)
graph.outputs.remove(old_var)
graph.outputs.insert(i, new_var)
def replace_variable(graph: Graph, old_var: Variable, new_var: Variable, with_assert: bool = True):
old_var.replace(new_var, with_assert=with_assert)
if old_var in graph.inputs:
i = graph.inputs.index(old_var)
graph.inputs.remove(old_var)
graph.inputs.insert(i, new_var)
if old_var in graph.outputs:
i = graph.outputs.index(old_var)
graph.outputs.remove(old_var)
graph.outputs.insert(i, new_var)
def main(k, s, p, n, h1, w1, c1, expected_shape_dict: Dict[Axis, int]):
orders = [OrderNHWC, OrderHWNC, OrderHWCN, OrderNCHW, OrderCNHW, OrderCHWN]
for order_x in orders:
op = MaxPooling2D(None, ksize=k, stride=s, padding=p)
x = Variable((n, h1, w1, c1), OrderNHWC)
x.change_order(order_x)
y, = op(x)
for axis in y.order.axes:
assert y.shape_dict[axis] == expected_shape_dict[axis]
def test_sgemm_invalid_C_shape():
op = Sgemm(None,
M=10,
N=20,
K=30,
out_shape=[1, 2, 3, 4],
out_order=OrderNHWC,
transpose_A=True,
transpose_B=True)
x = Variable((10, 30), OrderNC)
w = Variable((20, 30), OrderNC)
op(x, w)
def _replace_output(op: Operator, var_name: str,
target_orders: Union[Order, List[Order]]):
v = op.outputs[var_name]
if isinstance(target_orders, Order):
target_orders = [target_orders]
if v.order in target_orders:
return False
v_new = Variable(v.shape, v.order).change_order(target_orders[0])
op.replace_output(v, v_new, with_assert=False)
v_new.transpose(v.order).replace(v, with_assert=False)
return True
def test_every_order():
orders = [OrderNHWC, OrderHWNC, OrderHWCN, OrderNCHW, OrderCNHW, OrderCHWN]
for order in orders:
op = LocalResponseNormalization(None, n=1, k=2, alpha=0.1, beta=0.2)
x = Variable(np.arange(order.ndim) + 1, OrderNHWC)
x.change_order(order)
y, = op(x)
for axis in y.order.axes:
assert y.shape_dict[axis] == x.shape_dict[axis]
def template(x1_order=OrderNHWC,
x2_order=OrderNHWC,
y_order=OrderNHWC,
description: str = ""):
vx1 = np.random.rand(2, 3, 4, 5) - 0.5
vx2 = np.random.rand(2, 3, 4, 5) - 0.5
vy = vx1 / vx2
x1 = Variable(vx1.shape, order=OrderNHWC)
x2 = Variable(vx2.shape, order=OrderNHWC)
y = x1 / x2
x1.change_order(x1_order)
x2.change_order(x2_order)
y.change_order(y_order)
generate_kernel_test_case(
description=f"ElementwiseDiv {description}",
graph=Graph([x1, x2], [y]),
inputs={
x1: np.transpose(vx1,
[OrderNHWC.axes_dict[a] for a in x1.order.axes]),
x2: np.transpose(vx2,
[OrderNHWC.axes_dict[a] for a in x2.order.axes])
},
expected={
y: np.transpose(vy, [OrderNHWC.axes_dict[a] for a in y.order.axes])
},
)
def template(x_order=OrderNHWC, y_order=OrderNCHW, description: str = ""):
vx = np.random.rand(2, 3, 4, 5)
vy = np.transpose(vx, [x_order.axes_dict[a] for a in y_order.axes])
x = Variable(vx.shape, order=x_order)
y = x.transpose(y_order)
generate_kernel_test_case(
description=f"Transpose {description}",
backend=["webgpu", "webgl", "webassembly"],
graph=Graph([x], [y]),
inputs={x: vx},
expected={y: vy},
)
def exec(self):
x = self.inputs["x"]
x_shape_dict = x.shape_dict
N = x_shape_dict[Axis.N]
H2 = self.parameters["outsize"][0]
W2 = self.parameters["outsize"][1]
C2 = x_shape_dict[Axis.C]
y = Variable([N, H2, W2, C2], OrderNHWC)
y.change_order(
x.order
) # output same order as input to preserve following reshape semantics
self.append_output("y", y)
return y,
def template(shape=(2, 3, 4, 5),
x1_order=OrderNHWC,
x2_order=OrderNHWC,
y_order=OrderNHWC,
description: str = ""):
# vx1 = np.random.rand(*shape).astype(np.float32) - 0.5
# vx2 = np.random.rand(*shape).astype(np.float32) - 0.5
vx1 = np.arange(np.prod(shape)).reshape(shape).astype(np.float32)
vx2 = np.arange(np.prod(shape)).reshape(shape).astype(np.float32)
vy = vx1 + vx2
x1 = Variable(vx1.shape, order=OrderNHWC)
x2 = Variable(vx2.shape, order=OrderNHWC)
y = x1 + x2
x1.change_order(x1_order)
x2.change_order(x2_order)
y.change_order(y_order)
generate_kernel_test_case(
description=f"ElementwiseAdd {description}",
graph=Graph([x1, x2], [y]),
inputs={
x1: np.transpose(vx1,
[OrderNHWC.axes_dict[a] for a in x1.order.axes]),
x2: np.transpose(vx2,
[OrderNHWC.axes_dict[a] for a in x2.order.axes])
},
expected={
y: np.transpose(vy, [OrderNHWC.axes_dict[a] for a in y.order.axes])
},
)
def test_general():
for condition_custom in [
{},
{"x1_order": OrderNCHW, "x2_order": OrderHWCN}
]:
condition = dict(condition_default)
condition.update(condition_custom)
vx1 = np.random.rand(2, 3, 4, 5)
vx2 = np.random.rand(2, 3, 4, 5)
vy = vx1 ** vx2
x1 = Variable(vx1.shape, order=OrderNHWC)
x2 = Variable(vx2.shape, order=OrderNHWC)
y = x1 ** x2
x1.change_order(condition["x1_order"])
x2.change_order(condition["x2_order"])
y.change_order(condition["y_order"])
generate_kernel_test_case(
description=f"ElementwisePow: " + (", ".join([f"{k}={v}" for k, v in condition_custom.items()])),
backend=condition["backend"],
graph=Graph([x1, x2], [y]),
inputs={
x1: ConstantVariable(vx1, OrderNHWC).change_order(x1.order).data,
x2: ConstantVariable(vx2, OrderNHWC).change_order(x2.order).data
},
expected={y: ConstantVariable(vy, OrderNHWC).change_order(y.order).data},
raise_skip=False
)
raise SkipTest