Safe cosine similarity

ops_util.go

// CosineSimilarity ....
// lhs -> left-hand side
// rhs -> right-hand side
func CosineSimilarity(lhs, rhs *Node) *Node {
  g := lhs.Graph()
  dtype := lhs.DType()
  axis := -1  // Axis over which to calculate the cosine.

  // Mask for rows that are fully zero, for which cosine similary is not normally defined.
  lhsMask := ReduceAndKeep(IsZero(lhs), ReduceLogicalAnd, axis), 
  rhsMask := ReduceAndKeep(IsZero(rhs), ReduceLogicalAnd, axis)

  // Recover original shape, by broadcasting the mask where we just reduced.
  lhsMask = BroadcastToShape(lhs, lhs.Shape())
  rhsMask = BroadcastToShape(rhs, rhs.Shape())

  // Replace rows with all zeroes (lhsMask/rhsMask) with 1.
  // Any positive numerical safe number would work, since the final computation for 
  // those rows won't be used, as long as they are not NaNs.
  one := ScalarOne(g, dtype)
  lhs = Where(lhsMask, lhs, one)
  rhs = Where(lhsMask, rhs, one)

  // ... calculate similarity as usual

  // Arbitrarily set the similarity of the zero-rows (lhsMask or rhsMask) to zero.
  zero := ScalarZero(g, dtype)
  similarity = Where(LogicalOr(lhsMask, rhsMask), zero, similarity)
}