⚡️ Speed up function compute_float8_scale by 6%

keras/src/ops/numpy.py

@@ -3468,7 +3468,8 @@ def isfinite(x):
     Returns:
         Output boolean tensor.
     """
-    if any_symbolic_tensors((x,)):
+    # Fast path: only package for symbolic guard if necessary.
+    if any_symbolic_tensors(args=(x,)):
         return Isfinite().symbolic_call(x)
     return backend.numpy.isfinite(x)
 
@@ -5003,7 +5004,7 @@ def reciprocal(x):
     Returns:
         Output tensor, element-wise reciprocal of `x`.
     """
-    if any_symbolic_tensors((x,)):
+    if any_symbolic_tensors(args=(x,)):
         return Reciprocal().symbolic_call(x)
     return backend.numpy.reciprocal(x)
 
@@ -6211,12 +6212,12 @@ def where(condition, x1=None, x2=None):
         A tensor with elements from `x1` where `condition` is `True`, and
         elements from `x2` where `condition` is `False`.
     """
-    if (x1 is None and x2 is not None) or (x1 is not None and x2 is None):
+    if (x1 is None) != (x2 is None):
         raise ValueError(
             "`x1` and `x2` either both should be `None`"
             " or both should have non-None value."
         )
-    if any_symbolic_tensors((condition, x1, x2)):
+    if any_symbolic_tensors(args=(condition, x1, x2)):
         return Where().symbolic_call(condition, x1, x2)
     return backend.numpy.where(condition, x1, x2)
 
@@ -6323,7 +6324,7 @@ def divide(x1, x2):
     Returns:
         Output tensor, the quotient `x1/x2`, element-wise.
     """
-    if any_symbolic_tensors((x1, x2)):
+    if any_symbolic_tensors(args=(x1, x2)):
         return Divide().symbolic_call(x1, x2)
     return backend.numpy.divide(x1, x2)
 

keras/src/quantizers/quantizers.py

@@ -339,15 +339,21 @@ def grad(*args, upstream=None):
 
 @keras_export("keras.quantizers.compute_float8_scale")
 def compute_float8_scale(amax, scale, dtype_max, margin=0):
-    # The algorithm for computing the new scale is sourced from
-    # https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/jax.html#transformer_engine.jax.update_fp8_metas
-    # wherein the `original_scale` corresponds to the reciprocal of the
-    # `scale` passed in this function.
-    scale = ops.reciprocal(scale)
-    sf = ops.divide(ops.divide(dtype_max, amax), 2**margin)
-    sf = ops.where(amax > 0.0, sf, scale)
-    sf = ops.where(ops.isfinite(amax), sf, scale)
-    return ops.reciprocal(sf)
+    # Fast symbolic check for all involved arguments.
+    if any_symbolic_tensors(args=(amax, scale, dtype_max)):
+        scale_inv = ops.reciprocal(scale)
+        sf = ops.divide(ops.divide(dtype_max, amax), 2 ** margin)
+        sf = ops.where(amax > 0.0, sf, scale_inv)
+        sf = ops.where(ops.isfinite(amax), sf, scale_inv)
+        return ops.reciprocal(sf)
+    # Fast path: do all math with backend NumPy, as much as possible in single lines to reduce Python stack overhead.
+    scale_inv = backend.numpy.reciprocal(scale)
+    amax_finite = backend.numpy.isfinite(amax)
+    sf0 = backend.numpy.divide(dtype_max, amax)
+    sf1 = backend.numpy.divide(sf0, 2 ** margin)
+    sf2 = backend.numpy.where(amax > 0.0, sf1, scale_inv)
+    sf3 = backend.numpy.where(amax_finite, sf2, scale_inv)
+    return backend.numpy.reciprocal(sf3)
 
 
 @keras_export("keras.quantizers.compute_float8_amax_history")