(input_index, a)
| 451 | # 2. in input 0, broadcast axes come next, followed by summed axes |
| 452 | # 3. in input 1, summed axes come next, followed by broadcast axes |
| 453 | def sort_key(input_index, a): |
| 454 | if a in preserved_axes: |
| 455 | return (-1, a) |
| 456 | elif ((input_index == 0 and a in broadcast_axes[0]) or |
| 457 | (input_index == 1 and a in axes_to_sum)): |
| 458 | return (0, a) |
| 459 | else: |
| 460 | return (1, a) |
| 461 | |
| 462 | axis_labels = [t0_axis_labels, t1_axis_labels] |
| 463 | sorted_axes = [ |