FLOPs computation for softmax op. For softmax(input): equation: flops = 3 * (numel)total number of elements in the input tensor.
(input_shapes, attrs)
| 350 | |
| 351 | @register_flops("softmax") |
| 352 | def _softmax_flops(input_shapes, attrs): |
| 353 | """FLOPs computation for softmax op. |
| 354 | For softmax(input): |
| 355 | equation: flops = 3 * (numel)total number of elements in the input tensor. |
| 356 | """ |
| 357 | input = input_shapes.get('X')[0] |
| 358 | return prod(input) * 3 |
| 359 | |
| 360 | |
| 361 | @register_flops("transpose2") |