parameters: (graph, node)
return_statement: return ops.OpStats("flops", (k * output_count * 2))
docstring: Calculates the compute resources needed for BatchMatMul.
docstring_summary: Calculates the compute resources needed for BatchMatMul.
docstring_tokens: Calculates the compute resources needed for BatchMatMul .
function: def _calc_batch_mat_mul_flops(graph, node):
    """Calculates the compute resources needed for BatchMatMul."""
    transpose_a = node.attr["transpose_a"].b
    a_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
    a_shape.assert_is_fully_defined()
    if transpose_a:
      k = int(a_shape[-2])
    else:
      k = int(a_shape[-1])
    output_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
    output_shape.assert_is_fully_defined()
    output_count = np.prod(output_shape.as_list())
    return ops.OpStats("flops", (k * output_count * 2))
function_tokens: def _calc_batch_mat_mul_flops ( graph , node ) : transpose_a = node . attr [ "transpose_a" ] . b a_shape = graph_util . tensor_shape_from_node_def_name ( graph , node . input [ 0 ] ) a_shape . assert_is_fully_defined ( ) if transpose_a : k = int ( a_shape [ - 2 ] ) else : k = int ( a_shape [ - 1 ] ) output_shape = graph_util . tensor_shape_from_node_def_name ( graph , node . name ) output_shape . assert_is_fully_defined ( ) output_count = np . prod ( output_shape . as_list ( ) ) return ops . OpStats ( "flops" , ( k * output_count * 2 ) )