|
3 | 3 | import onnx
|
4 | 4 | import onnxscript as ost
|
5 | 5 | from onnxscript import opset19 as op # opset19 is the lastest by 202309
|
| 6 | +from onnxscript import opset11 |
| 7 | +from onnxscript import opset13 |
6 | 8 |
|
7 | 9 | np.random.seed(0)
|
8 | 10 |
|
@@ -69,8 +71,6 @@ def greater_input_dtype_int64(x: ost.FLOAT[27, 9]) ->ost.BOOL[27, 9]:
|
69 | 71 | return y
|
70 | 72 | make_model_and_data(greater_input_dtype_int64, np.random.randint(0, 100, size=[27, 9], dtype=np.int64), force_saving_input_as_dtype_float32=True, force_saving_output_as_dtype_float32=True)
|
71 | 73 |
|
72 |
| -from onnxscript import opset11 |
73 |
| - |
74 | 74 | @ost.script()
|
75 | 75 | def two_resizes_with_shared_subgraphs(x: ost.FLOAT["batch", 1, "height", "width"], y: ost.FLOAT[1, 1, 3, 2], z: ost.FLOAT[1, 1, 2, 1]) ->ost.FLOAT["batch", 1, "height", "width"]:
|
76 | 76 | shape_src_1 = opset11.Shape(x)
|
@@ -310,3 +310,32 @@ def einsum_const_inputs(input_0: ost.FLOAT[3, 2, 2, 4]) -> ost.FLOAT[3, 2, 2, 2]
|
310 | 310 | return op.Einsum(input_0, input_1, equation="bhwc, hkc -> bhwk")
|
311 | 311 |
|
312 | 312 | make_model_and_data(einsum_const_inputs, input_0_data)
|
| 313 | + |
| 314 | +''' This subgraph looks the same as LayerNorm expanded, but it has |
| 315 | + axes=1 in ReduceMean which does not meet the requirement of LayerNorm: |
| 316 | + - axes[-1] = -1 or the axis of last dimension |
| 317 | + - adjacent axes, e.g. [1, 2, 3] or [-3, -2, -1] |
| 318 | +''' |
| 319 | + |
| 320 | +n = 1 |
| 321 | +c = 4 |
| 322 | +h = w = 8 |
| 323 | +mul_weight = np.random.rand(c, 1, 1).astype(np.float32) |
| 324 | +add_weight = np.random.rand(c, 1, 1).astype(np.float32) |
| 325 | + |
| 326 | +@ost.script() |
| 327 | +def layer_norm_no_fusion(x: ost.FLOAT[n, c, h, w]) -> ost.FLOAT[n, c, h, w]: |
| 328 | + reduce_mean = opset13.ReduceMean(x, axes=[1], keepdims=1) |
| 329 | + sub = opset13.Sub(x, reduce_mean) |
| 330 | + |
| 331 | + pow = opset13.Pow(sub, opset13.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [], np.array([2], dtype=np.float32)))) |
| 332 | + reduce_mean_1 = opset13.ReduceMean(pow, axes=[1], keepdims=1) |
| 333 | + add = opset13.Add(reduce_mean_1, opset13.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [], np.array([9.999999974752427e-7], dtype=np.float32)))) |
| 334 | + sqrt = opset13.Sqrt(add) |
| 335 | + |
| 336 | + div = opset13.Div(sub, sqrt) |
| 337 | + mul = opset13.Mul(opset13.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [c, 1, 1], mul_weight)), div) |
| 338 | + add = opset13.Add(mul, opset13.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [c, 1, 1], add_weight))) |
| 339 | + |
| 340 | + return add |
| 341 | +make_model_and_data(layer_norm_no_fusion, np.random.rand(n, c, h, w).astype(np.float32)) |
0 commit comments