Luke Hutton | 261b7b6 | 2023-01-10 14:50:31 +0000 | [diff] [blame] | 1 | # Copyright (c) 2020-2023, ARM Limited. |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 2 | # SPDX-License-Identifier: Apache-2.0 |
Luke Hutton | 261b7b6 | 2023-01-10 14:50:31 +0000 | [diff] [blame] | 3 | import math |
| 4 | |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 5 | import numpy as np |
Jerry Ge | cb7201e | 2023-06-01 22:45:26 +0000 | [diff] [blame] | 6 | from tosa.DType import DType |
| 7 | |
| 8 | DTYPE_ATTRIBUTES = { |
| 9 | DType.BOOL: {"str": "b", "width": 1}, |
| 10 | DType.INT4: {"str": "i4", "width": 4}, |
| 11 | DType.INT8: {"str": "i8", "width": 8}, |
| 12 | DType.UINT8: {"str": "u8", "width": 8}, |
| 13 | DType.INT16: {"str": "i16", "width": 16}, |
| 14 | DType.UINT16: {"str": "u16", "width": 16}, |
| 15 | DType.INT32: {"str": "i32", "width": 32}, |
| 16 | DType.INT48: {"str": "i48", "width": 48}, |
| 17 | DType.FP16: {"str": "f16", "width": 16}, |
| 18 | DType.BF16: {"str": "bf16", "width": 16}, |
| 19 | DType.FP32: {"str": "f32", "width": 32}, |
| 20 | } |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 21 | |
| 22 | |
| 23 | class ArgGen: |
| 24 | """Argument generator functions. These functions take a shape and dtype to |
| 25 | create arguments for an operator. Methods are prefixed with 'ag' to make |
| 26 | search easy.""" |
| 27 | |
| 28 | def __init__(self): |
| 29 | pass |
| 30 | |
Jerry Ge | cb7201e | 2023-06-01 22:45:26 +0000 | [diff] [blame] | 31 | def typeStr(dtype): |
| 32 | if dtype in DTYPE_ATTRIBUTES: |
| 33 | return DTYPE_ATTRIBUTES[dtype]["str"] |
| 34 | else: |
| 35 | raise Exception("Unknown dtype, cannot convert to string: {}".format(dtype)) |
| 36 | |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 37 | @staticmethod |
| 38 | def agNone(op, shapes, rng): |
| 39 | """A trivial argument generator for operators that only take tensor |
| 40 | operands""" |
| 41 | return [("", [])] |
| 42 | |
| 43 | # Build the axis argument for operators where we want to iterate over N axes |
| 44 | # as an argument |
| 45 | @staticmethod |
| 46 | def agAxes(op, shapes, rng): |
| 47 | axes = [] |
Won Jeon | f9c0cee | 2023-09-18 16:32:45 -0700 | [diff] [blame^] | 48 | if shapes == (): |
| 49 | axes.append(["_axis_0", [0]]) |
| 50 | return axes |
| 51 | |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 52 | for i in range(-len(shapes), len(shapes), 1): |
| 53 | if i >= 0: |
| 54 | axes.append(["_axis_{}".format(i), [i]]) |
| 55 | else: |
| 56 | axes.append(["_axis_m{}".format(-i), [i]]) |
| 57 | return axes |
| 58 | |
| 59 | # Build the axis LIST argument for operators that take an axis list. |
| 60 | # This builds a list of each axis individually, plus one element |
| 61 | # that contains a list of all axes. Note that we need to pack the list in |
| 62 | # an additional list so that it isn't exploded when being passed to the |
| 63 | # build_operator function. |
| 64 | # tensor_arg_count not used |
| 65 | def agAxesList(op, shapes, rng): |
| 66 | axes = ArgGen.agAxes(op, shapes, rng) |
| 67 | axes_list = [] |
| 68 | for desc, a in axes: |
| 69 | axes_list.append([desc, [a]]) |
| 70 | |
| 71 | axes_list.append(["_axisall", [list(range(len(shapes)))]]) |
| 72 | axes_list.append(["_axisall_none", [None]]) |
| 73 | return axes_list |
| 74 | |
| 75 | def agAxesListKeepdims(op, shapes, rng): |
| 76 | axes = ArgGen.agAxes(op, shapes, rng) |
| 77 | axes_list = [] |
| 78 | for desc, a in axes: |
| 79 | axes_list.append([desc + "_keep0", [a, False]]) |
| 80 | # avoid trying to reduce an axis of shape 1, as the TFL converter |
| 81 | # will optimize away the entire reduction |
| 82 | if (a[0] >= 0 and shapes[a[0]] != 1) or ( |
| 83 | a[0] < 0 and shapes[len(shapes) + a[0]] != 1 |
| 84 | ): |
| 85 | axes_list.append([desc + "_keep1", [a, True]]) |
| 86 | |
| 87 | axes_list.append(["_axisall_keep0", [list(range(len(shapes))), False]]) |
| 88 | axes_list.append(["_axisall_keep0_none", [None, False]]) |
| 89 | # another instance where the reduce gets optimized out. |
| 90 | if len(shapes) != 1: |
| 91 | axes_list.append(["_axisall_keep1", [list(range(len(shapes))), True]]) |
| 92 | axes_list.append(["_axisall_keep1_none", [None, True]]) |
| 93 | # no longer test axis empty, as TFL converter optimizes the reduce out |
| 94 | return axes_list |
| 95 | |
| 96 | # conv2d argument generators build the TF constants |
| 97 | def agConv2d(op, shapes, rng): |
| 98 | arg_list = [] |
| 99 | |
| 100 | # Must be rank 4 |
| 101 | if len(shapes) < 4: |
| 102 | return arg_list |
| 103 | |
| 104 | filter_h, filter_w = op["filter"] |
| 105 | |
| 106 | # strides, padding, dilations, |
| 107 | for stride_h in [1, 2]: |
| 108 | for stride_w in [1, 2]: |
| 109 | for padding in ["SAME", "VALID"]: |
| 110 | for dilation_h in [1, 2]: |
| 111 | for dilation_w in [1, 2]: |
| 112 | |
| 113 | # Disqualify argument combinations that would cause |
| 114 | # an illegal convolution |
| 115 | |
| 116 | if (padding == "VALID") and ( |
| 117 | (shapes[1] - (filter_h - 1) * 2 - dilation_h) <= 0 |
| 118 | or (shapes[2] - (filter_w - 1) * 2 - dilation_w) <= 0 |
| 119 | ): |
| 120 | continue |
| 121 | |
Jeremy Johnson | 0e6218e | 2022-05-05 17:08:04 +0100 | [diff] [blame] | 122 | if ( |
| 123 | (shapes[1] - 1 - (filter_h - 1) * dilation_h) % stride_h |
| 124 | != 0 |
| 125 | ) or ( |
| 126 | (shapes[2] - 1 - (filter_w - 1) * dilation_w) % stride_w |
| 127 | != 0 |
| 128 | ): |
| 129 | # Not an exact integer output |
| 130 | continue |
| 131 | |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 132 | arg_list.append( |
| 133 | [ |
| 134 | "_st{}{}_pad{}_dilat{}{}".format( |
| 135 | stride_h, |
| 136 | stride_w, |
| 137 | padding, |
| 138 | dilation_h, |
| 139 | dilation_w, |
| 140 | ), |
| 141 | [ |
| 142 | [stride_h, stride_w], |
| 143 | padding, |
| 144 | [dilation_h, dilation_w], |
| 145 | ], |
| 146 | ] |
| 147 | ) |
| 148 | return arg_list |
| 149 | |
TatWai Chong | fd62905 | 2022-07-25 04:01:58 +0000 | [diff] [blame] | 150 | # conv3d argument generators build the TF constants |
| 151 | def agConv3d(op, shapes, rng): |
| 152 | arg_list = [] |
| 153 | |
| 154 | # input shape = [OC, KD, KH, KW, IC] |
| 155 | # Must be rank 5 |
| 156 | if len(shapes) != 5: |
| 157 | return arg_list |
| 158 | |
| 159 | if len(op["filter"]) < 3: |
| 160 | return arg_list |
| 161 | |
| 162 | filter_d, filter_h, filter_w = op["filter"] |
| 163 | |
| 164 | # strides, padding, dilations, |
| 165 | for stride_d in [1, 2]: |
| 166 | for stride_h in [1, 2]: |
| 167 | for stride_w in [1, 2]: |
| 168 | for padding in ["SAME", "VALID"]: |
| 169 | for dilation_d in [1, 2]: |
| 170 | for dilation_h in [1, 2]: |
| 171 | for dilation_w in [1, 2]: |
| 172 | |
| 173 | # Disqualify argument combinations that would cause |
| 174 | # an illegal convolution |
| 175 | # fmt: off |
| 176 | if (padding == "VALID") and ( |
| 177 | (shapes[1] - (filter_d - 1) * 2 - dilation_d) <= 0 |
| 178 | or (shapes[2] - (filter_h - 1) * 2 - dilation_h) <= 0 |
| 179 | or (shapes[3] - (filter_w - 1) * 2 - dilation_w) <= 0 |
| 180 | ): |
| 181 | continue |
| 182 | |
| 183 | if ( |
| 184 | (shapes[1] - 1 - (filter_d - 1) * dilation_d) % stride_d |
| 185 | != 0 |
| 186 | ) or ( |
| 187 | (shapes[2] - 1 - (filter_h - 1) * dilation_h) % stride_h |
| 188 | != 0 |
| 189 | ) or ( |
| 190 | (shapes[3] - 1 - (filter_w - 1) * dilation_w) % stride_w |
| 191 | != 0 |
| 192 | ): |
| 193 | # Not an exact integer output |
| 194 | continue |
| 195 | # fmt: on |
| 196 | |
| 197 | # TODO investigate the error of `CPU implementation of Conv3D |
| 198 | # currently only supports dilated rates of 1.` from Tensorflow. |
| 199 | # Only test dilations = [1, 1, 1, 1, 1] for now. |
| 200 | if ( |
| 201 | (dilation_d != 1) |
| 202 | or (dilation_h != 1) |
| 203 | or (dilation_w != 1) |
| 204 | ): |
| 205 | continue |
| 206 | |
| 207 | # Tensorflow expects strides is a list of ints that has length >= 5. |
| 208 | # Strides and dilations in the batch and depth dimensions must be 1. |
| 209 | arg_list.append( |
| 210 | [ |
| 211 | "_st{}{}{}{}{}_pad{}_dilat{}{}{}{}{}".format( |
| 212 | 1, |
| 213 | stride_d, |
| 214 | stride_h, |
| 215 | stride_w, |
| 216 | 1, |
| 217 | padding, |
| 218 | 1, |
| 219 | dilation_d, |
| 220 | dilation_h, |
| 221 | dilation_w, |
| 222 | 1, |
| 223 | ), |
| 224 | [ |
| 225 | [1, stride_d, stride_h, stride_w, 1], |
| 226 | padding, |
| 227 | [ |
| 228 | 1, |
| 229 | dilation_d, |
| 230 | dilation_h, |
| 231 | dilation_w, |
| 232 | 1, |
| 233 | ], |
| 234 | ], |
| 235 | ] |
| 236 | ) |
| 237 | return arg_list |
| 238 | |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 239 | # conv2d argument generators build the TF constants |
| 240 | def agDepthwiseConv2d(op, shapes, rng): |
| 241 | arg_list = [] |
| 242 | |
| 243 | # Must be rank 4 |
| 244 | if len(shapes) < 4: |
| 245 | return arg_list |
| 246 | |
| 247 | filter_h, filter_w = op["filter"] |
| 248 | |
| 249 | # strides, padding, dilations, Depthwise conv2d is the same as conv2d |
| 250 | # except that strides in h/w must be the same and the argument must be |
| 251 | # formatted as [1, stride_h, stride_w, 1] in TF. |
| 252 | for stride in [1, 2]: |
| 253 | for padding in ["SAME", "VALID"]: |
| 254 | for dilation_h in [1, 2]: |
| 255 | for dilation_w in [1, 2]: |
| 256 | |
| 257 | # Disqualify argument combinations that would cause an illegal |
| 258 | # convolution |
| 259 | |
| 260 | if (padding == "VALID") and ( |
| 261 | (shapes[1] - (filter_h - 1) * 2 - dilation_h) <= 0 |
| 262 | or (shapes[2] - (filter_w - 1) * 2 - dilation_w) <= 0 |
| 263 | ): |
| 264 | continue |
| 265 | |
| 266 | # When dilation is used, stride must be 1x1 (TF rules) |
| 267 | if dilation_h > 1 or dilation_w > 1: |
| 268 | if stride > 1: |
| 269 | continue |
| 270 | |
| 271 | # Dilation must evenly divide the tensor. Some of our inputs |
| 272 | # intentionally use odd-sized tensors. |
| 273 | if shapes[1] % dilation_h != 0 or shapes[2] % dilation_w != 0: |
| 274 | continue |
| 275 | |
Jeremy Johnson | 0e6218e | 2022-05-05 17:08:04 +0100 | [diff] [blame] | 276 | if ( |
| 277 | (shapes[1] - 1 - (filter_h - 1) * dilation_h) % stride != 0 |
| 278 | ) or ( |
| 279 | (shapes[2] - 1 - (filter_w - 1) * dilation_w) % stride != 0 |
| 280 | ): |
| 281 | # Not an exact integer output |
| 282 | continue |
| 283 | |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 284 | arg_list.append( |
| 285 | [ |
| 286 | "_st{}{}_pad{}_dilat{}{}".format( |
| 287 | stride, stride, padding, dilation_h, dilation_w |
| 288 | ), |
| 289 | [ |
| 290 | [1, stride, stride, 1], |
| 291 | padding, |
| 292 | [dilation_h, dilation_w], |
| 293 | ], |
| 294 | ] |
| 295 | ) |
| 296 | return arg_list |
| 297 | |
| 298 | # conv2d argument generators build the TF constants |
| 299 | def agTransposeConv2d(op, shapes, rng): |
| 300 | arg_list = [] |
| 301 | |
| 302 | # Must be rank 4 |
| 303 | if len(shapes) < 4: |
| 304 | return arg_list |
| 305 | |
| 306 | filter_h, filter_w = op["filter"] |
| 307 | |
| 308 | # strides, padding, dilations, |
| 309 | for stride_h in [1, 2]: |
| 310 | for stride_w in [1, 2]: |
| 311 | for padding in ["SAME", "VALID"]: |
| 312 | if padding == "SAME": |
| 313 | out_height = (shapes[1]) * stride_h |
| 314 | out_width = (shapes[2]) * stride_w |
| 315 | else: # padding == 'VALID' |
| 316 | out_height = (shapes[1] - 1) * stride_h + filter_h |
| 317 | out_width = (shapes[2] - 1) * stride_w + filter_w |
| 318 | |
| 319 | output_shape = [shapes[0], out_height, out_width, shapes[3] * 2] |
| 320 | arg_list.append( |
| 321 | [ |
| 322 | "_st{}{}_pad{}".format(stride_h, stride_w, padding), |
| 323 | [output_shape, [stride_h, stride_w], padding], |
| 324 | ] |
| 325 | ) |
| 326 | return arg_list |
| 327 | |
| 328 | def agPooling(op, shapes, rng): |
| 329 | arg_list = [] |
| 330 | |
| 331 | # Must be rank 4 |
| 332 | if len(shapes) < 4: |
| 333 | return arg_list |
| 334 | |
| 335 | for stride_h in [1, 2]: |
| 336 | for stride_w in [1, 2]: |
| 337 | for kernel_h in [1, 2]: |
| 338 | for kernel_w in [1, 2]: |
| 339 | for padding in ["SAME", "VALID"]: |
| 340 | |
| 341 | if (padding == "VALID") and ( |
| 342 | (shapes[1] % (kernel_h * stride_h) > 0) |
| 343 | or (shapes[2] % (kernel_w * stride_w) > 0) |
| 344 | or (shapes[1] <= kernel_h) |
| 345 | or (shapes[2] <= kernel_w) |
| 346 | ): |
| 347 | continue |
| 348 | |
| 349 | if (padding == "SAME") and ( |
| 350 | (shapes[1] < kernel_h) or (shapes[2] < kernel_w) |
| 351 | ): |
| 352 | continue |
| 353 | |
Jeremy Johnson | 0e6218e | 2022-05-05 17:08:04 +0100 | [diff] [blame] | 354 | if ((shapes[1] - kernel_h) % stride_h != 0) or ( |
| 355 | (shapes[2] - kernel_w) % stride_w != 0 |
| 356 | ): |
| 357 | # Not an exact integer output |
| 358 | continue |
| 359 | |
Jerry Ge | cb7201e | 2023-06-01 22:45:26 +0000 | [diff] [blame] | 360 | # Note: tf.nn.avg_pool2d API doesn't support setting accumtype |
| 361 | # setting a dummy value to the test name as an reminder |
| 362 | accum_dtype = ArgGen.typeStr(DType.INT32) |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 363 | arg_list.append( |
| 364 | [ |
Jerry Ge | cb7201e | 2023-06-01 22:45:26 +0000 | [diff] [blame] | 365 | "_st{}{}_pad{}_kern{}{}_acc{}".format( |
| 366 | stride_h, |
| 367 | stride_w, |
| 368 | padding, |
| 369 | kernel_h, |
| 370 | kernel_w, |
| 371 | accum_dtype, |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 372 | ), |
| 373 | [ |
| 374 | [stride_h, stride_w], |
| 375 | [kernel_h, kernel_w], |
| 376 | padding, |
| 377 | ], |
| 378 | ] |
| 379 | ) |
| 380 | return arg_list |
| 381 | |
| 382 | def getFactors(val, start=1): |
| 383 | factors = [] |
| 384 | for i in range(start, int(np.sqrt(val))): |
| 385 | if (val % i) == 0: |
| 386 | factors.append(i) |
| 387 | |
| 388 | return factors |
| 389 | |
| 390 | def agReshape(op, shapes, rng): |
| 391 | # This is slow code. Fortunately, the numbers involved are small |
| 392 | arg_list = [] |
| 393 | |
| 394 | total_elements = 1 |
| 395 | for s in shapes: |
| 396 | total_elements *= s |
| 397 | |
| 398 | # Find integer factors of this shape |
| 399 | factors = ArgGen.getFactors(total_elements) |
| 400 | |
| 401 | for rank in range(1, len(shapes) + 1): |
| 402 | if len(factors) < rank: |
| 403 | break |
| 404 | |
| 405 | new_shape = [] |
| 406 | remaining_elements = total_elements |
| 407 | |
| 408 | # Randomly shuffle the factors and iteratively pick from the factors |
| 409 | # of the remaining elements |
| 410 | shuffled_factors = rng.permutation(factors) |
| 411 | for i in range(rank): |
| 412 | # Pick rank - 1 factors |
| 413 | new_shape.append(shuffled_factors[0]) |
| 414 | remaining_elements = remaining_elements // shuffled_factors[0] |
| 415 | shuffled_factors = rng.permutation( |
| 416 | ArgGen.getFactors(remaining_elements) |
| 417 | ) |
| 418 | new_shape.append(remaining_elements) |
| 419 | |
| 420 | # Don't do no-op reshapes because TFLite optimizes out the op |
| 421 | if new_shape == list(shapes): |
| 422 | continue |
| 423 | |
| 424 | arg_list.append(["_rank{}".format(rank), [new_shape]]) |
| 425 | |
| 426 | return arg_list |
| 427 | |
| 428 | def agTranspose(op, shapes, rng): |
| 429 | arg_list = [] |
| 430 | |
| 431 | # Must have at least two dimensions to transpose |
| 432 | if (len(shapes)) < 2: |
| 433 | return arg_list |
| 434 | |
| 435 | # Pick a bunch of random permutations |
| 436 | range_arr = np.arange(len(shapes)) |
| 437 | for i in range(len(shapes)): |
| 438 | perm = rng.permutation(range_arr).astype(np.int32) |
| 439 | # print('\n shape {} permute{} perm: {} arr: {}'.format(shapes, i, |
| 440 | # perm, range_arr)) |
| 441 | if np.allclose(perm, range_arr): |
| 442 | print("skipped") |
| 443 | continue |
| 444 | arg_list.append(["_permute{}".format(i), [perm]]) |
| 445 | |
| 446 | return arg_list |
| 447 | |
| 448 | def agSlice(op, shapes, rng): |
| 449 | arg_list = [] |
| 450 | |
| 451 | rank = len(shapes) |
| 452 | |
| 453 | if rank == 1 and shapes[0] == 1: |
| 454 | return arg_list |
| 455 | |
| 456 | for i in range(4): |
| 457 | # Pick a few random start points, axes, and strides |
| 458 | start = np.empty((rank), dtype=int) |
| 459 | size = np.empty((rank), dtype=int) |
| 460 | for j in range(rank): |
| 461 | if shapes[j] > 2: |
| 462 | start[j] = rng.integers(0, shapes[j] - 2) |
| 463 | # print('j = {}: {} - {} - 1: {}'.format(j, shapes[j], |
| 464 | # start[j], shapes[j] - start[j] - 1)) |
| 465 | size[j] = rng.integers(1, shapes[j] - start[j] - 1) |
| 466 | else: |
| 467 | start[j] = 0 |
| 468 | size[j] = shapes[j] |
| 469 | |
| 470 | arg_list.append(["_perm{}".format(i), [start, size]]) |
| 471 | |
| 472 | return arg_list |
| 473 | |
| 474 | def agStridedSlice(op, shapes, rng): |
| 475 | arg_list = [] |
| 476 | |
| 477 | rank = len(shapes) |
| 478 | |
| 479 | # Reference model is limited to rank=6 internally right now |
| 480 | if rank > 3: |
| 481 | return arg_list |
| 482 | |
| 483 | if rank == 1 and shapes[0] == 1: |
| 484 | return arg_list |
| 485 | |
| 486 | for i in range(4): |
| 487 | # Pick a few random begin points, axes, and strides |
| 488 | begin = np.empty((rank), dtype=int) |
| 489 | end = np.empty((rank), dtype=int) |
| 490 | strides = np.empty((rank), dtype=int) |
| 491 | |
| 492 | begin_mask = rng.integers(0, (1 << (rank - 1))) |
| 493 | end_mask = rng.integers(0, (1 << (rank - 1))) |
| 494 | |
| 495 | for j in range(rank): |
| 496 | |
| 497 | if begin_mask & (1 << j) or shapes[j] < 2: |
| 498 | begin[j] = 0 |
| 499 | else: |
| 500 | begin[j] = rng.integers(0, shapes[j] - 1) |
| 501 | |
| 502 | if end_mask & (1 << j) or shapes[j] < 2 or (begin[j] + 2) >= shapes[j]: |
| 503 | end[j] = shapes[j] |
| 504 | else: |
| 505 | end[j] = rng.integers(begin[j] + 1, shapes[j] - 1) |
| 506 | |
| 507 | possible_stride = ArgGen.getFactors(end[j] - begin[j], 2) |
| 508 | |
| 509 | if not possible_stride: |
| 510 | strides[j] = 1 |
| 511 | else: |
| 512 | strides[j] = rng.choice(possible_stride) |
| 513 | |
| 514 | # Randomly set the masks, except ellipsis_mask and new_axis_mask |
| 515 | # which must be zero for now For begin/end mask this to work, |
| 516 | # strides must be adjusted to still be divsible... |
| 517 | ellipsis_mask = 0 |
| 518 | new_axis_mask = 0 |
| 519 | |
| 520 | # if rng.choice([0, 1]) and rank > 1: |
| 521 | # new_axis_mask = 1 << rng.integers(0, rank - 1) |
| 522 | # else: |
| 523 | # new_axis_mask = 0 |
| 524 | |
| 525 | if rng.choice([0, 1]) and rank > 1: |
| 526 | shrink_axis_mask = 1 << rng.integers(0, rank - 1) |
| 527 | else: |
| 528 | shrink_axis_mask = 0 |
| 529 | |
| 530 | # Only one of these bits may be set. Prefer shrink_axis_mask |
| 531 | new_axis_mask = new_axis_mask & ~shrink_axis_mask |
| 532 | |
| 533 | arg_list.append( |
| 534 | [ |
| 535 | "_perm{}".format(i), |
| 536 | [ |
| 537 | begin, |
| 538 | end, |
| 539 | strides, |
| 540 | begin_mask, |
| 541 | end_mask, |
| 542 | ellipsis_mask, |
| 543 | new_axis_mask, |
| 544 | shrink_axis_mask, |
| 545 | ], |
| 546 | ] |
| 547 | ) |
| 548 | |
| 549 | # print('Shape: {} begin={} end={} strides={} begin_mask={:x} |
| 550 | # end_mask={:x} new_axis_mask={:x} shrink_mask={:x}'.format(shapes, |
| 551 | # begin, end, strides, begin_mask, end_mask, new_axis_mask, |
| 552 | # shrink_axis_mask)) |
| 553 | |
| 554 | return arg_list |
| 555 | |
| 556 | # tf.stack axis can be [0, rank(input)] |
| 557 | def agStack(op, shapes, rng): |
| 558 | axes = [] |
| 559 | for i in range(len(shapes) + 1): |
| 560 | axes.append(["_axis{}".format(i), [i]]) |
| 561 | return axes |
| 562 | |
TatWai Chong | f7008da | 2022-09-09 09:35:40 +0000 | [diff] [blame] | 563 | def agMirrorPad(op, shapes, rng): |
| 564 | arg_list = [] |
| 565 | |
| 566 | rank = len(shapes) |
| 567 | for mode in ["REFLECT", "SYMMETRIC"]: |
| 568 | for left in range(3): |
| 569 | for right in range(3): |
| 570 | paddings = np.zeros((rank, 2), dtype=np.int32) |
| 571 | is_valid = True |
| 572 | |
| 573 | # Fill in the padding parameter if the values are valid on each dimension, |
| 574 | # otherwise drop that case. |
| 575 | for d in range(rank): |
| 576 | paddings[d, 0] = left |
| 577 | paddings[d, 1] = right |
| 578 | |
| 579 | # In "REFLECT" mode, paddings must be no greater than tensor dim size - 1. |
| 580 | if mode == "REFLECT": |
| 581 | if (left > shapes[d] - 1) or (right > shapes[d] - 1): |
| 582 | is_valid = False |
| 583 | break |
| 584 | |
| 585 | # In "SYMMETRIC" mode, paddings must be no greater than tensor dim size. |
| 586 | else: |
| 587 | if (left > shapes[d]) or (right > shapes[d]): |
| 588 | is_valid = False |
| 589 | break |
| 590 | |
| 591 | if is_valid: |
| 592 | arg_list.append( |
| 593 | [ |
| 594 | "_pad{}{}_{}".format(left, right, mode[0:3].lower()), |
| 595 | [paddings, mode], |
| 596 | ] |
| 597 | ) |
| 598 | return arg_list |
| 599 | |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 600 | def agPad(op, shapes, rng): |
| 601 | arg_list = [] |
| 602 | |
| 603 | rank = len(shapes) |
| 604 | for left in range(3): |
| 605 | for right in range(3): |
TatWai Chong | 2226f90 | 2023-02-22 18:38:01 -0800 | [diff] [blame] | 606 | # Padding nothing in tensorflow lite causes the interpreter fail to set |
| 607 | # the input tensor properly due to date type mismatch. |
| 608 | if (left == 0) and (right == 0): |
| 609 | continue |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 610 | |
TatWai Chong | 2226f90 | 2023-02-22 18:38:01 -0800 | [diff] [blame] | 611 | # A simple way to generate explicit pad_const including zero. |
| 612 | pad_const = (left - right) * rng.integers(0, 5, dtype=np.int32) |
| 613 | padding = np.zeros((rank, 2), dtype=np.int32) |
| 614 | for d in range(rank): |
| 615 | padding[d, 0] = left |
| 616 | padding[d, 1] = right |
| 617 | |
| 618 | arg_list.append( |
| 619 | ["_pad{}{}".format(left, right), [padding, pad_const]] |
| 620 | ) |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 621 | return arg_list |
| 622 | |
TatWai Chong | 0cef07e | 2023-02-27 13:22:52 -0800 | [diff] [blame] | 623 | def agResize(op, shapes, rng): |
| 624 | args = [] |
| 625 | for mode in ["nearest", "bilinear"]: |
| 626 | for align_corners in [True, False]: |
| 627 | for half_pixel in [True, False]: |
| 628 | # If half_pixel_centers is True, align_corners must be False. |
Jerry Ge | c937609 | 2023-03-08 17:50:49 +0000 | [diff] [blame] | 629 | if (align_corners is True) and (half_pixel is True): |
TatWai Chong | 0cef07e | 2023-02-27 13:22:52 -0800 | [diff] [blame] | 630 | continue |
| 631 | |
| 632 | for i in range(1, 4): |
| 633 | args.append( |
| 634 | [ |
| 635 | "_{}_align{}_half{}_scale{}".format( |
| 636 | mode, int(align_corners), int(half_pixel), i |
| 637 | ), |
| 638 | [mode, align_corners, half_pixel, i], |
| 639 | ] |
| 640 | ) |
| 641 | return args |
| 642 | |
Jeremy Johnson | 015c355 | 2022-02-23 12:15:03 +0000 | [diff] [blame] | 643 | def agFill(op, shapes, rng): |
| 644 | values = [] |
| 645 | for i in range(4): |
| 646 | value = rng.integers(0, 10, dtype=np.int32) |
| 647 | values.append(["_value{}".format(value), [shapes, value]]) |
| 648 | return values |
| 649 | |
| 650 | def getValuesToSum(total, rng): |
| 651 | # Get a list of random integers that sum up to 'total' |
| 652 | vals = [] |
| 653 | |
| 654 | # np.random.randint() min and max to be different, so if the remainder |
| 655 | # is 1, give up |
| 656 | while total > 1: |
| 657 | vals.append(rng.integers(1, total)) |
| 658 | total = total - vals[-1] |
| 659 | |
| 660 | if total == 1: |
| 661 | vals.append(1) |
| 662 | |
| 663 | return vals |
| 664 | |
| 665 | def agSplit(op, shapes, rng): |
| 666 | arg_list = [] |
| 667 | |
| 668 | rank = len(shapes) |
| 669 | |
| 670 | # Shuffle the random number generator a few more times to get |
| 671 | # a better range of axes across shapes |
| 672 | for i in range(rank): |
| 673 | for j in range(shapes[i]): |
| 674 | rng.integers(shapes[i]) |
| 675 | |
| 676 | for i in range(3): |
| 677 | # Need to generate tests for both the num_splits and size_vector versions. |
| 678 | axis = rng.choice(np.arange(0, rank)) |
| 679 | |
| 680 | # For num_splits, get a few divisors of the given axis |
| 681 | divs = ArgGen.getFactors(shapes[axis], 2) |
| 682 | |
| 683 | if divs: |
| 684 | # Get no more than 2 samples |
| 685 | splits = list(rng.choice(divs, size=2)) |
| 686 | |
| 687 | for s in splits: |
| 688 | arg_list.append( |
| 689 | ["_split{}_axis{}".format(int(s), axis), [int(s), axis]] |
| 690 | ) |
| 691 | |
| 692 | # For vector splits, get a list of integers that sum up to the axis size |
| 693 | vals = ArgGen.getValuesToSum(shapes[axis], rng) |
| 694 | |
| 695 | if len(vals) > 1: |
| 696 | arg_list.append(["_splitv_axis{}".format(axis), [vals, axis]]) |
| 697 | |
| 698 | return arg_list |
| 699 | |
| 700 | def agTile(op, shapes, rng): |
| 701 | arg_list = [] |
| 702 | |
| 703 | rank = len(shapes) |
| 704 | |
| 705 | # create 1D multiples list |
| 706 | multiples = list() |
| 707 | for i in range(rank): |
| 708 | multiples.append(rng.integers(1, 4)) |
| 709 | |
| 710 | multiples_str = "x".join(list(str(i) for i in multiples)) |
| 711 | |
| 712 | arg_list.append(["_tile_{}".format(multiples_str), [multiples]]) |
| 713 | |
| 714 | return arg_list |
| 715 | |
| 716 | def agGather(op, shapes, rng): |
| 717 | args = [] |
| 718 | for batch_dims in range(len(shapes) - 1): |
| 719 | for axis in range(batch_dims, len(shapes)): |
| 720 | # indices value must be within [0, shapes[i]) |
| 721 | |
| 722 | # Create an arbitrary shape for the indices |
| 723 | indices_rank = rng.integers(batch_dims + 1, 4) |
| 724 | indices_shape = rng.integers(1, 8, size=indices_rank) |
| 725 | |
| 726 | # Copy in the batch dimensions because they must match |
| 727 | for b in range(batch_dims): |
| 728 | indices_shape[b] = shapes[b] |
| 729 | |
| 730 | # Calculate total element count |
| 731 | indices_size = 1 |
| 732 | for j in range(indices_rank): |
| 733 | indices_size = indices_shape[j] * indices_size |
| 734 | |
| 735 | indices = rng.integers(0, shapes[axis], indices_size, np.int32).reshape( |
| 736 | indices_shape |
| 737 | ) |
| 738 | |
| 739 | args.append( |
| 740 | [ |
| 741 | "_batchdims_{}_axis_{}".format(batch_dims, axis), |
| 742 | [indices, batch_dims, axis], |
| 743 | ] |
| 744 | ) |
| 745 | return args |
| 746 | |
| 747 | def agGatherND(op, shapes, rng): |
| 748 | args = [] |
| 749 | |
| 750 | for N in range(1, len(shapes) - 1): |
| 751 | # Rank includes the N dimension |
| 752 | indices_rank = rng.integers(2, 4, size=1)[0] |
| 753 | indices_shape = [] |
| 754 | |
| 755 | indices_shape = rng.integers(1, 8, size=indices_rank) |
| 756 | indices_shape[-1] = N |
| 757 | |
| 758 | indices_count = 1 |
| 759 | for i in range(indices_rank - 1): |
| 760 | indices_count = indices_count * indices_shape[i] |
| 761 | |
| 762 | indices_list = np.zeros(shape=(indices_count, N), dtype=np.int32) |
| 763 | |
| 764 | for i in range(indices_count): |
| 765 | for j in range(N): |
| 766 | indices_list[i, j] = rng.integers(0, shapes[j], size=1)[0] |
| 767 | |
| 768 | indices = indices_list.reshape(indices_shape) |
| 769 | |
| 770 | args.append(["_n{}".format(N), [indices]]) |
| 771 | |
| 772 | return args |
| 773 | |
| 774 | def agScatterND(op, shapes, rng): |
| 775 | args = [] |
| 776 | |
| 777 | # ScatterND has to generate a constant shapes tensor, indices |
| 778 | # tensor, and a tensor of updates. Unforunately, the updates |
| 779 | # need to be a size that's based on the N generated in this |
| 780 | # function and the dtype known only in the TensorGen function, |
| 781 | # but not in ArgGen. |
| 782 | # |
| 783 | # There are many bad ways to solve this and we'll choose the |
| 784 | # least of the evils which still gives reasonable coverage of |
| 785 | # the possible operand shapes. |
| 786 | for N in range(1, len(shapes)): |
| 787 | # Rank includes the N dimension |
| 788 | indices_rank = rng.integers(2, 4, size=1)[0] |
| 789 | indices_shape = [] |
| 790 | |
| 791 | indices_shape = rng.integers(1, 8, size=indices_rank) |
| 792 | indices_shape[-1] = N |
| 793 | |
| 794 | # Store the Shapes, and the indicies value tensor as arguments. |
| 795 | args.append(["_n{}".format(N), [shapes, indices_shape, N, rng]]) |
| 796 | |
| 797 | return args |
| 798 | |
| 799 | def agSpaceToBatch(op, shapes, rng): |
| 800 | batch_rank = 1 |
| 801 | channel_rank = 1 |
| 802 | block_rank = len(shapes) - batch_rank - channel_rank |
| 803 | |
| 804 | # must have at least rank 1 (M) block |
| 805 | if block_rank < 1: |
| 806 | return [] |
| 807 | |
| 808 | args = [] |
| 809 | block_shape = [] |
| 810 | padding_shape = [] |
| 811 | |
| 812 | for i in range(block_rank): |
| 813 | block_size = 2 |
| 814 | padding_size = block_size - (shapes[i + 1] % block_size) |
| 815 | block_shape.append(block_size) |
| 816 | padding_shape.append([0, padding_size]) |
| 817 | |
| 818 | args.append(["_blockrank_{}".format(block_rank), [block_shape, padding_shape]]) |
| 819 | return args |
| 820 | |
| 821 | def agBatchToSpace(op, shapes, rng): |
| 822 | batch_rank = 1 |
| 823 | channel_rank = 1 |
| 824 | block_rank = len(shapes) - batch_rank - channel_rank |
| 825 | |
| 826 | # must have at least rank 1 (M) block |
| 827 | if block_rank < 1: |
| 828 | return [] |
| 829 | |
| 830 | args = [] |
| 831 | block_shape = [] |
| 832 | padding_shape = [] |
| 833 | block_prod = 1 |
| 834 | |
| 835 | for i in range(block_rank): |
| 836 | block_size = 2 |
| 837 | block_prod = block_prod * block_size |
| 838 | crop_size = 0 |
| 839 | block_shape.append(block_size) |
| 840 | padding_shape.append([0, crop_size]) |
| 841 | |
| 842 | # batch / prod(block_shape[i]) must be integer |
| 843 | # transpose to swap depth and batch. so shape[-1] would be batch dim |
| 844 | if shapes[-1] % block_prod == 0: |
| 845 | args.append( |
| 846 | ["_blockrank_{}".format(block_rank), [block_shape, padding_shape]] |
| 847 | ) |
| 848 | |
| 849 | return args |
| 850 | |
| 851 | def agSpaceToDepth(op, shapes, rng): |
| 852 | # must be rank 4 input tensor |
| 853 | if len(shapes) != 4: |
| 854 | return [] |
| 855 | |
| 856 | block_size = 2 |
| 857 | |
| 858 | # spatial dimension must be divisible by block_size |
| 859 | if shapes[1] % block_size != 0 or shapes[2] % block_size != 0: |
| 860 | return [] |
| 861 | |
| 862 | args = [] |
| 863 | args.append(["_blocksize_{}".format(block_size), [block_size]]) |
| 864 | |
| 865 | return args |
| 866 | |
| 867 | def agDepthToSpace(op, shapes, rng): |
| 868 | # must be rank 4 input tensor |
| 869 | if len(shapes) != 4: |
| 870 | return [] |
| 871 | |
| 872 | block_size = 2 |
| 873 | # depth dimension must be divisible by block_size * block_size |
| 874 | if shapes[3] % (block_size * block_size) != 0: |
| 875 | return [] |
| 876 | |
| 877 | args = [] |
| 878 | args.append(["_blocksize_{}".format(block_size), [block_size]]) |
| 879 | |
| 880 | return args |
| 881 | |
| 882 | def agFakequant(op, shapes, rng): |
| 883 | args = [] |
| 884 | for num_bits in [8, 16]: |
| 885 | for narrow in [False, True]: |
| 886 | args.append( |
| 887 | ["_bits{}_narrow{}".format(num_bits, narrow), [num_bits, narrow]] |
| 888 | ) |
| 889 | |
| 890 | return args |
| 891 | |
| 892 | def agShift(op, shapes, rng): |
| 893 | args = [] |
| 894 | |
| 895 | for shift in rng.integers(0, 32, size=8): |
| 896 | args.append(["_shift{}".format(shift), [shift]]) |
| 897 | |
| 898 | return args |
| 899 | |
| 900 | def agFloat(op, shapes, rng): |
| 901 | args = [] |
| 902 | |
| 903 | i = 0 |
| 904 | for alpha in np.float32(rng.random(size=2)): |
| 905 | args.append(["_{}".format(i), [alpha]]) |
| 906 | |
| 907 | return args |
| 908 | |
| 909 | # Similar to agAxes, but tf.OneHot only allow axis from [-1, rank(input)] |
| 910 | def agOneHot(op, shapes, rng): |
| 911 | axes = [] |
| 912 | for i in range(-1, len(shapes) + 1, 1): |
| 913 | if i >= 0: |
| 914 | axes.append(["_axis_{}".format(i), [i]]) |
| 915 | else: |
| 916 | axes.append(["_axis_m{}".format(-i), [i]]) |
| 917 | return axes |
Luke Hutton | 261b7b6 | 2023-01-10 14:50:31 +0000 | [diff] [blame] | 918 | |
| 919 | def agRFFT2d(op, shape, rng): |
| 920 | args = [] |
| 921 | |
| 922 | # Must be rank 3 input tensor |
| 923 | if len(shape) != 3: |
| 924 | return [] |
| 925 | |
| 926 | # Check rfft2d with enforced fft_length |
| 927 | for fft_length_h in [2, 32]: |
| 928 | for fft_length_w in [2, 8, 16]: |
| 929 | fft_length = [fft_length_h, fft_length_w] |
| 930 | args.append(["_fft_length_{}x{}".format(*fft_length), [fft_length]]) |
| 931 | |
| 932 | # Check rfft2d with no fft_length provided (fft_length=None). |
| 933 | # In this case, the height and width of the input should be |
| 934 | # used for the calculation. Therefore, we need to check that |
| 935 | # the input shape is already a power of two. |
| 936 | def is_power_of_two(x): |
| 937 | return math.log(x, 2).is_integer() |
| 938 | |
| 939 | height, width = shape[1:3] |
| 940 | if is_power_of_two(height) and is_power_of_two(width): |
| 941 | args.append(["_fft_length_None", [None]]) |
| 942 | |
| 943 | return args |