def forward(
self,
lhs_pos: FloatTensorType,
rhs_pos: FloatTensorType,
lhs_neg: FloatTensorType,
rhs_neg: FloatTensorType,
) -> Tuple[FloatTensorType, FloatTensorType, FloatTensorType]:
num_chunks, num_pos_per_chunk, dim = match_shape(lhs_pos, -1, -1, -1)
match_shape(rhs_pos, num_chunks, num_pos_per_chunk, dim)
match_shape(lhs_neg, num_chunks, -1, dim)
match_shape(rhs_neg, num_chunks, -1, dim)
# Smaller distances are higher scores, so take their negatives.
pos_scores = (
(lhs_pos.float() - rhs_pos.float())
.pow_(2)
.sum(dim=-1)
.clamp_min_(1e-30)
.sqrt_()
.neg()
)
lhs_neg_scores = batched_all_pairs_l2_dist(rhs_pos, lhs_neg).neg()
rhs_neg_scores = batched_all_pairs_l2_dist(lhs_pos, rhs_neg).neg()
return pos_scores, lhs_neg_scores, rhs_neg_scores
def forward(
self,
lhs_pos: FloatTensorType,
rhs_pos: FloatTensorType,
lhs_neg: FloatTensorType,
rhs_neg: FloatTensorType,
) -> Tuple[FloatTensorType, FloatTensorType, FloatTensorType]:
num_chunks, num_pos_per_chunk, dim = match_shape(lhs_pos, -1, -1, -1)
match_shape(rhs_pos, num_chunks, num_pos_per_chunk, dim)
match_shape(lhs_neg, num_chunks, -1, dim)
match_shape(rhs_neg, num_chunks, -1, dim)
# Equivalent to (but faster than) torch.einsum('cid,cid->ci', ...).
pos_scores = (lhs_pos.float() * rhs_pos.float()).sum(-1)
# Equivalent to (but faster than) torch.einsum('cid,cjd->cij', ...).
lhs_neg_scores = torch.bmm(rhs_pos, lhs_neg.transpose(-1, -2))
rhs_neg_scores = torch.bmm(lhs_pos, rhs_neg.transpose(-1, -2))
return pos_scores, lhs_neg_scores, rhs_neg_scores