(pfor_input)
| 2108 | |
| 2109 | @RegisterPFor("MatMul") |
| 2110 | def _convert_matmul(pfor_input): |
| 2111 | # TODO(agarwal): Check if tiling is faster than two transposes. |
| 2112 | a, a_stacked, _ = pfor_input.input(0) |
| 2113 | b, b_stacked, _ = pfor_input.input(1) |
| 2114 | tr_a = pfor_input.get_attr("transpose_a") |
| 2115 | tr_b = pfor_input.get_attr("transpose_b") |
| 2116 | if a_stacked and b_stacked: |
| 2117 | output = wrap(math_ops.matmul(a, b, adjoint_a=tr_a, adjoint_b=tr_b), True) |
| 2118 | return output |
| 2119 | elif a_stacked: |
| 2120 | if tr_a: |
| 2121 | a = array_ops.transpose(a, [0, 2, 1]) |
| 2122 | if a.shape.is_fully_defined(): |
| 2123 | x, y, z = a.shape |
| 2124 | else: |
| 2125 | x, y, z = [ |
| 2126 | array_ops.reshape(i, []) |
| 2127 | for i in array_ops.split(array_ops.shape(a), 3) |
| 2128 | ] |
| 2129 | a = array_ops.reshape(a, [x * y, z]) |
| 2130 | prod = math_ops.matmul(a, b, transpose_b=tr_b) |
| 2131 | return wrap(array_ops.reshape(prod, [x, y, -1]), True) |
| 2132 | else: |
| 2133 | assert b_stacked |
| 2134 | if tr_b: |
| 2135 | perm = [2, 0, 1] |
| 2136 | b = array_ops.transpose(b, perm) |
| 2137 | else: |
| 2138 | # As an optimization, if one of the first two dimensions is 1, then we can |
| 2139 | # reshape instead of transpose. |
| 2140 | # TODO(agarwal): This check can be done inside Transpose kernel. |
| 2141 | b_shape = array_ops.shape(b) |
| 2142 | min_dim = math_ops.minimum(b_shape[0], b_shape[1]) |
| 2143 | perm = control_flow_ops.cond( |
| 2144 | math_ops.equal(min_dim, 1), lambda: [0, 1, 2], lambda: [1, 0, 2]) |
| 2145 | new_shape = array_ops.stack([b_shape[1], b_shape[0], b_shape[2]]) |
| 2146 | b = array_ops.transpose(b, perm) |
| 2147 | b = array_ops.reshape(b, new_shape) |
| 2148 | |
| 2149 | if b.shape.is_fully_defined(): |
| 2150 | x, y, z = b.shape |
| 2151 | else: |
| 2152 | x, y, z = [ |
| 2153 | array_ops.reshape(i, []) |
| 2154 | for i in array_ops.split(array_ops.shape(b), 3) |
| 2155 | ] |
| 2156 | b = array_ops.reshape(b, [x, y * z]) |
| 2157 | prod = math_ops.matmul(a, b, transpose_a=tr_a) |
| 2158 | prod = array_ops.reshape(prod, [-1, y, z]) |
| 2159 | prod = array_ops.transpose(prod, [1, 0, 2]) |
| 2160 | return wrap(prod, True) |
| 2161 | |
| 2162 | |
| 2163 | # TODO(rmlarsen): Use the converter of BatchMatMulV2 once compatibility window |
nothing calls this directly
no test coverage detected