MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / _convert_batch_mat_mul

Function _convert_batch_mat_mul

tensorflow/python/ops/parallel_for/pfor.py:2166–2179  ·  view source on GitHub ↗
(pfor_input)

Source from the content-addressed store, hash-verified

2164# is met.
2165@RegisterPFor("BatchMatMul")
2166def _convert_batch_mat_mul(pfor_input):
2167 # TODO(agarwal): There may be a more efficient way to do this instead of
2168 # stacking the inputs.
2169 pfor_input.stack_inputs()
2170 x = pfor_input.stacked_input(0)
2171 y = pfor_input.stacked_input(1)
2172 adj_x = pfor_input.get_attr("adj_x")
2173 adj_y = pfor_input.get_attr("adj_y")
2174
2175 x = _flatten_first_two_dims(x)
2176 y = _flatten_first_two_dims(y)
2177 output = math_ops.matmul(x, y, adjoint_a=adj_x, adjoint_b=adj_y)
2178 output = _unflatten_first_dim(output, pfor_input.pfor.loop_len_vector)
2179 return wrap(output, True)
2180
2181
2182@RegisterPFor("BatchMatMulV2")

Callers

nothing calls this directly

Calls 7

_unflatten_first_dimFunction · 0.85
stack_inputsMethod · 0.80
stacked_inputMethod · 0.80
_flatten_first_two_dimsFunction · 0.70
wrapFunction · 0.70
get_attrMethod · 0.45
matmulMethod · 0.45

Tested by

no test coverage detected