Merge pull request jax-ml#25117 from pearu:pearu/arcsin-mac-arm

PiperOrigin-RevId: 702012134

Author: Google-ML-Automation

jax/_src/test_util.py

@@ -25,6 +25,7 @@
 import logging
 import math
 import os
+import platform
 import re
 import sys
 import tempfile
@@ -1704,6 +1705,10 @@ def complex_plane_sample(dtype, size_re=10, size_im=None):
     size_im = size_re
   finfo = np.finfo(dtype)
 
+  machine = platform.machine()
+  is_arm_cpu = machine.startswith('aarch') or machine.startswith('arm')
+  smallest = np.nextafter(finfo.tiny, finfo.max) if is_arm_cpu and platform.system() == 'Darwin' else finfo.tiny
+
   def make_axis_points(size):
     prec_dps_ratio = 3.3219280948873626
     logmin = logmax = finfo.maxexp / prec_dps_ratio
@@ -1722,8 +1727,8 @@ def make_axis_points(size):
       axis_points[1] = finfo.min
       axis_points[-2] = finfo.max
     if size > 0:
-      axis_points[size] = -finfo.tiny
-      axis_points[-size - 1] = finfo.tiny
+      axis_points[size] = -smallest
+      axis_points[-size - 1] = smallest
     axis_points[0] = -np.inf
     axis_points[-1] = np.inf
     return axis_points

tests/lax_test.py

@@ -4400,34 +4400,14 @@ def regions_with_inaccuracies_keep(*to_keep):
     elif name == 'tanh':
       regions_with_inaccuracies_keep('ninf', 'pinf', 'ninfj', 'pinfj')
 
-    elif name == 'arcsin':
-      if is_arm_cpu and platform.system() == 'Darwin':
-        regions_with_inaccuracies_keep('q1.real', 'q2.real', 'q3.real', 'q4.real', 'neg.real', 'pos.real')
-      else:
-        regions_with_inaccuracies.clear()
-
-    elif name == 'arcsinh':
-      if is_arm_cpu and platform.system() == 'Darwin':
-        regions_with_inaccuracies_keep('q1.imag', 'q2.imag', 'q3.imag', 'q4.imag',
-                                       'negj.imag', 'posj.imag')
-      else:
-        regions_with_inaccuracies.clear()
-
     elif name == 'arccos':
       regions_with_inaccuracies_keep('q4.imag', 'ninf', 'pinf', 'ninfj', 'pinfj.real')
 
     elif name in {'cos', 'sin'}:
       regions_with_inaccuracies_keep('ninf.imag', 'pinf.imag')
 
-    elif name == 'log1p':
-      if is_arm_cpu and platform.system() == 'Darwin':
-        regions_with_inaccuracies_keep('q1.imag', 'q2.imag', 'q3.imag', 'q4.imag', 'negj.imag',
-                                       'posj.imag')
-      else:
-        regions_with_inaccuracies.clear()
-
-    elif name in {'positive', 'negative', 'conjugate', 'sin', 'cos', 'sqrt', 'expm1', 'tan',
-                  'arcsinh', 'arccosh', 'arctan', 'arctanh', 'square'}:
+    elif name in {'positive', 'negative', 'conjugate', 'sin', 'cos', 'sqrt', 'expm1', 'tan', 'log1p',
+                  'arcsin', 'arcsinh', 'arccosh', 'arctan', 'arctanh', 'square'}:
       regions_with_inaccuracies.clear()
     else:
       assert 0  # unreachable