title: Pytorch maximum inner product without materializing intermediate
body: <p>I'd like to do the equivalent of</p>
  <pre><code># X.shape = (batch, M, D)
  # Y.shape = (N, D)
  Z = torch.max(torch.bmm(X, Y.T), dim=2)
  </code></pre>
  <p>while only paying memory cost <code>(batch, M)</code> by not materializing an intermediate of size <code>(batch, M, N)</code> (because, as you may have guessed, I run out of GPU ram).  I can do this manually by slicing the bmm and accumulating the max, but that's alot of cuda kernel calls, so it kinda sucks.</p>
  <p>Is there a better solution?</p>