Make _quad_form_derivative for BlockDiagLazy batch-compatible
If the left_vec was passed in with an explicit batch dimension but no explicit trailing dimension things ended up breaking in adding batch dimensions (and broke in much less obvious cases if some implicit broadcasting was happening).
See the discussion on for details.
@@ -74,9 +74,12 @@ class BlockLazyTensor(LazyTensor):
return res
def _quad_form_derivative(self, left_vecs, right_vecs):
- if left_vecs.ndimension() == 1:
- left_vecs = left_vecs.unsqueeze(1)
- right_vecs = right_vecs.unsqueeze(1)
+ if left_vecs.ndim == 1:
+ left_vecs = left_vecs.unsqueeze(-1)
+ right_vecs = right_vecs.unsqueeze(-1)
+ # deal with left_vecs having batch dimensions
+ elif left_vecs.size(-1) != right_vecs.size(-1):
+ left_vecs = left_vecs.unsqueeze(-1)
left_vecs = self._add_batch_dim(left_vecs)
right_vecs = self._add_batch_dim(right_vecs)
res = self.base_lazy_tensor._quad_form_derivative(left_vecs, right_vecs)
