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_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
