Relax unit vector check for bfloat16

Author: thecoop

server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java

@@ -26,7 +26,8 @@
 
 /** Utility class for vector quantization calculations */
 public class BQVectorUtils {
-    private static final float EPSILON = 1e-4f;
+    // NOTE: this is currently > 1e-4f due to bfloat16
+    private static final float EPSILON = 1e-2f;
 
     public static double sqrtNewtonRaphson(double x, double curr, double prev) {
         return (curr == prev) ? curr : sqrtNewtonRaphson(x, 0.5 * (curr + x / curr), curr);

server/src/main/java/org/elasticsearch/index/codec/vectors/OptimizedScalarQuantizer.java

@@ -109,7 +109,7 @@ public QuantizationResult[] multiScalarQuantize(
     }
 
     public QuantizationResult scalarQuantize(float[] vector, float[] residualDestination, int[] destination, byte bits, float[] centroid) {
-        assert similarityFunction != COSINE || VectorUtil.isUnitVector(vector);
+        assert similarityFunction != COSINE || BQVectorUtils.isUnitVector(vector);
         assert similarityFunction != COSINE || VectorUtil.isUnitVector(centroid);
         assert vector.length <= destination.length;
         assert bits > 0 && bits <= 8;

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryFlatVectorsScorer.java

@@ -27,11 +27,13 @@
 import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
 import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer;
 import org.elasticsearch.index.codec.vectors.BQSpaceUtils;
+import org.elasticsearch.index.codec.vectors.BQVectorUtils;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
 import org.elasticsearch.simdvec.ESVectorUtil;
 
 import java.io.IOException;
 
+import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
 import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN;
 import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT;
 
@@ -69,6 +71,7 @@ public RandomVectorScorer getRandomVectorScorer(
             assert binarizedVectors.size() > 0 : "BinarizedByteVectorValues must have at least one vector for ES816BinaryFlatVectorsScorer";
             OptimizedScalarQuantizer quantizer = binarizedVectors.getQuantizer();
             float[] centroid = binarizedVectors.getCentroid();
+            assert similarityFunction != COSINE || BQVectorUtils.isUnitVector(target);
             float[] scratch = new float[vectorValues.dimension()];
             int[] initial = new int[target.length];
             byte[] quantized = new byte[BQSpaceUtils.B_QUERY * binarizedVectors.discretizedDimensions() / 8];