ai-edge-quantizer-nightly 0.4.0.dev20251015__py3-none-any.whl → 0.4.0.dev20251016__py3-none-any.whl

ai_edge_quantizer/utils/validation_utils.py

@@ -40,6 +40,8 @@ def get_validation_func(
     return median_diff_ratio
   elif func_name == "cosine_similarity":
     return cosine_similarity
+  elif func_name == "kl_divergence":
+    return kl_divergence
   else:
     raise ValueError(f"Validation function {func_name} not supported")
 
@@ -134,6 +136,43 @@ def cosine_similarity(
   return np.dot(data1, data2) / (norm_data1 * norm_data2)
 
 
+def kl_divergence(
+    data1: np._typing.ArrayLike,
+    data2: np._typing.ArrayLike,
+    epsilon: float = 1e-9,
+) -> float:
+  """Calculates the KL divergence between data1 & data2.
+
+  KL(data2 || data1) = sum(data2 * log(data2 / data1)).
+  data2 is treated as the true distribution P, and data1 as the
+  approximated distribution Q.
+  Non-positive values in data1 and data2 are clipped to 0 before
+  KL divergence calculation. Epsilon is added to avoid log(0) and
+  division by zero.
+
+  Args:
+    data1: input data to be used for comparison (distribution Q)
+    data2: input data to be used for comparison (distribution P),
+      data1 & 2 must be of the same shape
+    epsilon: small value to avoid log(0) and division by zero.
+
+  Returns:
+    A float value representing the KL divergence between data1 & 2.
+
+  Raises:
+    Value error if the two inputs don't have the same number of elements.
+  """
+  data1, data2 = _preprocess_same_size_arrays(data1, data2)
+  # special handling for tensor of size 0
+  if data1.size == 0:
+    return float(0)
+
+  p = np.maximum(0, data2)
+  q = np.maximum(0, data1)
+
+  return float(np.sum(p * np.log((p + epsilon) / (q + epsilon))))
+
+
 def _preprocess_same_size_arrays(
     data1: np._typing.ArrayLike, data2: np._typing.ArrayLike
 ) -> Tuple[np.ndarray, np.ndarray]:

ai_edge_quantizer/utils/validation_utils_test.py

@@ -106,6 +106,34 @@ class ValidationUtilTest(googletest.TestCase):
     result = validation_utils.cosine_similarity(data1, data2)
     self.assertEqual(result, 0)
 
+  def test_kl_divergence(self):
+    data1 = [0.5, 0.5]
+    data2 = [0.1, 0.9]
+    result = validation_utils.kl_divergence(data1, data2)
+    self.assertAlmostEqual(result, 0.36808, 4)
+
+  def test_kl_divergence_zero_in_q(self):
+    data1 = [0, 1]
+    data2 = [1, 0]
+    result = validation_utils.kl_divergence(data1, data2)
+    self.assertAlmostEqual(result, 20.7232658, 4)
+
+  def test_kl_divergence_negative_values(self):
+    data1 = [-1, 1]
+    data2 = [1, -1]
+    result = validation_utils.kl_divergence(data1, data2)
+    self.assertAlmostEqual(result, 20.7232658, 4)
+
+  def test_kl_divergence_0d(self):
+    data1 = []
+    data2 = []
+    result = validation_utils.kl_divergence(data1, data2)
+    self.assertEqual(result, 0)
+
+  def test_get_validation_func_kl_divergence(self):
+    func = validation_utils.get_validation_func("kl_divergence")
+    self.assertEqual(func, validation_utils.kl_divergence)
+
 
 if __name__ == "__main__":
   googletest.main()

{ai_edge_quantizer_nightly-0.4.0.dev20251015.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20251016.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-quantizer-nightly
-Version: 0.4.0.dev20251015
+Version: 0.4.0.dev20251016
 Summary: A quantizer for advanced developers to quantize converted AI Edge models.
 Home-page: https://github.com/google-ai-edge/ai-edge-quantizer
 Keywords: On-Device ML,AI,Google,TFLite,Quantization,LLMs,GenAI

{ai_edge_quantizer_nightly-0.4.0.dev20251015.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20251016.dist-info}/RECORD

@@ -72,10 +72,10 @@ ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=42OWzQsRTXq3XQYmoxlz177_d
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
 ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=EoVjI_hplX_Rml3hfRsGmQOihexmizeJqt4SQcET9aA,14925
 ai_edge_quantizer/utils/tfl_interpreter_utils_test.py,sha256=6fjkM-rycZ95L4yfvlr0TN6RlrhfPzxNUYrZaYO_F0A,12013
-ai_edge_quantizer/utils/validation_utils.py,sha256=0sOdH4pzk_Pwh1r8O47iaECRng1Xn0ABn9GVc8UPNcY,4994
-ai_edge_quantizer/utils/validation_utils_test.py,sha256=1sblJWHLTYTbn1Qi9rwnrREOSXRy5KwHAWSwgI1e_aU,3697
-ai_edge_quantizer_nightly-0.4.0.dev20251015.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20251015.dist-info/METADATA,sha256=aPMv8uCtui-_UA9TOyLpHkk5ZQtCBkYRa-_568uIth4,1508
-ai_edge_quantizer_nightly-0.4.0.dev20251015.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20251015.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20251015.dist-info/RECORD,,
+ai_edge_quantizer/utils/validation_utils.py,sha256=QTYyQ_HDVrFTGPIsrA240Lv8tUw1fwWp2fu9kTVISkE,6224
+ai_edge_quantizer/utils/validation_utils_test.py,sha256=lO51rGskhzpXePRdZMU87u_YO35_sDp9_eQ85CmupL4,4600
+ai_edge_quantizer_nightly-0.4.0.dev20251016.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20251016.dist-info/METADATA,sha256=Pb491ErUpqbWSgau8BYJwALz1jrD-ziX6af_1Jg_caE,1508
+ai_edge_quantizer_nightly-0.4.0.dev20251016.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20251016.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20251016.dist-info/RECORD,,