[TFL] Fix BatchMatMul constant RHS. #46749
Conversation
For constant RHS and !adj_y, transpose of RHS only occurs once, so the allocation should be persistent across runs.
|
Verified that this fixes the issue. |
| CreateBatchMatMulOptions(builder_, adj_x, adj_y).Union(), | ||
| matmul_inputs, matmul_outputs); | ||
|
|
||
| // This is just a dummy op to check if the values of transposed RHS are |
There was a problem hiding this comment.
Can you explain why the neg op help here?
There was a problem hiding this comment.
It's somehow tricky. When there is no following ops (not limited to neg), temporary allocation with kTfLiteArenaRw type tends to re-claim the same memory across each evaluation and no other ops will modify values at that memory address (cause no other memory allocations take place). Therefore, the values of rhs are not polluted without neg op, and next evaluation can also get the same piece of memory, so the results are correct.
https://colab.research.google.com/drive/1eFVszZZN-NcR64X_VHns0w-l3Vwci0CY?usp=sharing
Here is the colab link from original post. Values of last tensor (temp tensor, transposed rhs) become
[[-51. -84.]
[ 3. 6.]
[ 7. 9.]]
where -51, -84 are the last two elements of Identity (output of neg op). While without neg op, everything works smoothly.
There was a problem hiding this comment.
Acked. Thanks for explanation. May be it worths a comment then.
There was a problem hiding this comment.
Updated comments. See if it's better! Thanks!
google-ml-butler
bot
added
kokoro:force-run
google-ml-butler
bot
added
kokoro:force-run
thaink
added
ready to pull
Fixes #46724. For constant RHS and !adj_y, transpose of RHS only occurs once, so the allocation should be persistent across runs. Old commit does this, but the allocation is accidentally overridden to kTfLiteArenaRw for constant RHS.
tensorflow/tensorflow/lite/kernels/batch_matmul.cc
Lines 193 to 199 in 7c0a3f1
/cc @abattery for visibility.