Calculates the compute resources needed for Einsum.
(graph, node)
| 571 | |
| 572 | @ops.RegisterStatistics('Einsum', 'flops') |
| 573 | def _einsum_flops(graph, node): |
| 574 | """Calculates the compute resources needed for Einsum.""" |
| 575 | assert len(node.input) == 2 |
| 576 | x_shape = tf.compat.v1.graph_util.tensor_shape_from_node_def_name( |
| 577 | graph, node.input[0]) |
| 578 | y_shape = tf.compat.v1.graph_util.tensor_shape_from_node_def_name( |
| 579 | graph, node.input[1]) |
| 580 | x_shape.assert_is_fully_defined() |
| 581 | y_shape.assert_is_fully_defined() |
| 582 | x_shape = x_shape.as_list() |
| 583 | y_shape = y_shape.as_list() |
| 584 | equation = str(node.attr['equation']) |
| 585 | equation = ( |
| 586 | equation.replace('s:', '') |
| 587 | .replace('"', '') |
| 588 | .replace(' ', '') |
| 589 | .replace('\n', '') |
| 590 | ) |
| 591 | x_str = equation.split(',')[0] |
| 592 | y_r_str = equation.split(',')[1] |
| 593 | y_str = y_r_str.split('->')[0] |
| 594 | r_str = y_r_str.split('->')[1] |
| 595 | shape_dic = {} |
| 596 | contracted = set() |
| 597 | for indice in x_str + y_str: |
| 598 | if indice in x_str: |
| 599 | indice_dim = x_shape[x_str.find(indice)] |
| 600 | elif indice in y_str: |
| 601 | indice_dim = y_shape[y_str.find(indice)] |
| 602 | else: |
| 603 | raise ValueError('indice {} not found in inputs'.format(indice)) |
| 604 | shape_dic[indice] = indice_dim |
| 605 | if indice not in r_str: |
| 606 | contracted.add(indice) |
| 607 | madds = np.prod([shape_dic[indice] for indice in r_str]) * ( |
| 608 | np.prod([shape_dic[indice] for indice in contracted])) |
| 609 | flops = 2 * madds |
| 610 | return ops.OpStats('flops', flops) |