wip

jax_scaled_arithmetics/lax/scaled_ops_common.py

@@ -15,6 +15,7 @@
     Shape,
     as_scaled_array,
     get_scale_dtype,
+    is_static_anyscale,
     is_static_zero,
     safe_div,
 )
@@ -223,10 +224,10 @@ def scaled_le(lhs: ScaledArray, rhs: ScaledArray) -> Array:
 def scaled_minmax(prim: jax.core.Primitive, lhs: ScaledArray, rhs: ScaledArray) -> ScaledArray:
     """General min/max scaled translation: propagating the largest input scale."""
     check_scalar_scales(lhs, rhs)
-    # Specific rule if lhs/rhs is zero => propagate the other term scale.
-    if np.all(is_static_zero(lhs)):
+    # Specific rule if lhs/rhs is zero or inf => propagate the other term scale.
+    if np.all(is_static_anyscale(lhs)):
         return ScaledArray(prim.bind(lhs.data, rhs.data), rhs.scale)
-    if np.all(is_static_zero(rhs)):
+    if np.all(is_static_anyscale(rhs)):
         return ScaledArray(prim.bind(lhs.data, rhs.data), lhs.scale)
 
     # Power-of-2 stable!

tests/lax/test_scaled_ops_l2.py

@@ -158,9 +158,8 @@ def test__scaled_addsub__not_overflowing_scale(self, prim):
         assert np.isfinite(z.scale)
         npt.assert_array_almost_equal(z, prim.bind(np.asarray(x, np.float32), np.asarray(y, np.float32)), decimal=6)
 
-    @parameterized.parameters(
-        {"prim": lax.max_p},
-        {"prim": lax.min_p},
+    @parameterized.product(
+        prim=[lax.min_p, lax.max_p],
     )
     def test__scaled_minmax__static_zero_scale_propagation(self, prim):
         scaled_op, _ = find_registered_scaled_op(prim)
@@ -172,6 +171,19 @@ def test__scaled_minmax__static_zero_scale_propagation(self, prim):
         # Keep the lhs scale.
         npt.assert_almost_equal(z.scale, 4.0)
 
+    @parameterized.product(
+        prim=[lax.min_p, lax.max_p],
+    )
+    def test__scaled_minmax__static_inf_scale_propagation(self, prim):
+        scaled_op, _ = find_registered_scaled_op(prim)
+        x = scaled_array([-1.0, 2.0], 4.0, dtype=np.float32, npapi=np)
+        y = scaled_array([-np.inf, np.inf], np.inf, dtype=np.float32, npapi=np)
+        z = scaled_op(x, y)
+        assert isinstance(z, ScaledArray)
+        assert z.dtype == x.dtype
+        # Keep the lhs scale.
+        npt.assert_almost_equal(z.scale, 4.0)
+
     def test__scaled_mul__proper_scaling(self):
         x = scaled_array([-2.0, 2.0], 3, dtype=np.float32)
         y = scaled_array([1.5, 1.5], 2, dtype=np.float32)

tests/lax/test_scipy_integration.py

@@ -25,22 +25,16 @@ def fn(a):
     @chex.variants(with_jit=False, without_jit=True)
     @parameterized.parameters(
         {"dtype": np.float32},
-        # {"dtype": np.float16},
+        {"dtype": np.float16},
     )
     def test__scipy_logsumexp__accurate_scaled_op(self, dtype):
-        import jax
         from jax.scipy.special import logsumexp
 
         input_scaled = scaled_array(self.rs.rand(10), 4.0, dtype=dtype)
-
-        print(jax.make_jaxpr(logsumexp)(input_scaled.data))
-
         # JAX `logsumexp` Jaxpr is a non-trivial graph!
         out_scaled = self.variant(autoscale(logsumexp))(input_scaled)
         out_expected = logsumexp(np.asarray(input_scaled))
         assert out_scaled.dtype == out_expected.dtype
         # Proper accuracy + keep the same scale.
         npt.assert_array_equal(out_scaled.scale, input_scaled.scale)
         npt.assert_array_almost_equal(out_scaled, out_expected, decimal=5)
-
-        assert False