On GPU, test_kernels.py::TestMaternKernel::test_compute[with_jit-arrays_degree_greater_than_zero] fails with small numerical errors

[tp832944]

What's the problem?

On an NVIDIA GPU, test_kernels.py::TestMaternKernel::test_compute[with_jit-arrays_degree_greater_than_zero] is failing with what appear to be small numerical errors.

How can we reproduce the issue?

Run test_kernels.py::TestMaternKernel::test_compute[with_jit-arrays_degree_greater_than_zero]

Test fails

Python version

3.12

Package version

0.3.1

Operating system

Ubuntu 24.04 with NVIDIA GPU

Other packages

jax[cuda12]

Relevant log output

___ TestMaternKernel.test_compute[with_jit-arrays_degree_greater_than_zero] ____

self = <tests.unit.test_kernels.TestMaternKernel object at 0x7a0cbb58ae10>
jit_variant = <function filter_jit at 0x7a0d7b51b600>
problem = _Problem(x=array([[0, 1, 2, 3],
[5, 6, 7, 8]]), y=array([[1, 2, 3, 4],
[5, 6, 7, 8]]), expected_output=a...5.5047366e-07],
[1.4261089e-05, 9.9999952e-01]]), kernel=MaternKernel(length_scale=1, output_scale=1, degree=1))

def test_compute(
    self, jit_variant: Callable[[Callable], Callable], problem: _Problem
) -> None:
    """Test ``compute`` method of ``coreax.kernels.ScalarValuedKernel``."""
    x, y, expected_output, kernel = problem
    output = jit_variant(kernel.compute)(x, y)

  np.testing.assert_array_almost_equal(output, expected_output)

tests/unit/test_kernels.py:87:

---

/usr/lib/python3.12/contextlib.py:81: in inner
return func(*args, **kwds)
.venv/lib/python3.12/site-packages/numpy/_utils/init.py:85: in wrapper
return fun(*args, **kwargs)

---

args = (<function assert_array_almost_equal..compare at 0x7a0c7e155ee0>, Array([[1.3973139e-01, 5.5047417e-07],
...0e-05, 1.0000012e+00]], dtype=float32), array([[1.3973127e-01, 5.5047366e-07],
[1.4261089e-05, 9.9999952e-01]]))
kwds = {'err_msg': '', 'header': 'Arrays are not almost equal to 6 decimals', 'precision': 6, 'verbose': True}

@wraps(func)
def inner(*args, **kwds):
    with self._recreate_cm():

      return func(*args, **kwds)

E AssertionError:
E Arrays are not almost equal to 6 decimals
E
E Mismatched elements: 1 / 4 (25%)
E Max absolute difference among violations: 1.6720929e-06
E Max relative difference among violations: 1.6720937e-06
E ACTUAL: array([[1.397314e-01, 5.504742e-07],
E [1.426110e-05, 1.000001e+00]], dtype=float32)
E DESIRED: array([[1.397313e-01, 5.504737e-07],
E [1.426109e-05, 9.999995e-01]])

/usr/lib/python3.12/contextlib.py:81: AssertionError

[gw265981]

I replicated this and can confirm that setting the environment variable, export JAX_ENABLE_X64=True, to enable double precision fixes the issue. I will test the performance impact of this.

My feeling is that we shouldn't change the kernel code, but add a general warning somewhere in the documentation.