[tensor_transform] Fix confusing usage of transpose