Repository: scipy/scipy
File: scipy/stats/tests/test_correlation.py (Python)

Code before:
        res = stats.chatterjeexi(x, y, method=method,
                                 y_continuous=y_continuous, axis=-1)
        ref = stats.chatterjeexi(x, y, y_continuous=y_continuous, axis=-1)
        np.testing.assert_allclose(res.statistic, ref.statistic, rtol=1e-15)
        np.testing.assert_allclose(res.pvalue, ref.pvalue, rtol=2e-2)

    def test_input_validation(self):
        rng = np.random.default_rng(25932435798274926)
        x, y = rng.random(size=(2, 10))

        message = 'Array shapes are incompatible for broadcasting.'
        with pytest.raises(ValueError, match=message):
            stats.chatterjeexi(x, y[:-1])

        message = '...axis 10 is out of bounds for array...'
        with pytest.raises(ValueError, match=message):
            stats.chatterjeexi(x, y, axis=10)

        message = '`y_continuous` must be boolean.'
        with pytest.raises(ValueError, match=message):
            stats.chatterjeexi(x, y, y_continuous='a herring')

        message = "`method` must be 'asymptotic' or"
        with pytest.raises(ValueError, match=message):
            stats.chatterjeexi(x, y, method='ekki ekii')

    def test_special_cases(self):
        message = 'One or more sample arguments is too small...'
        with pytest.warns(SmallSampleWarning, match=message):
            res = stats.chatterjeexi([1], [2])

        assert np.isnan(res.statistic)
        assert np.isnan(res.pvalue)


Review comment (suggestion):
this looks good--testing `jax` is retained modulo the exception handling, which I think is not a surprising limitation per previous discussions

Code after:
        n = np.floor(rng.uniform(100, 150)).astype(int)
        shape = (2, n)
        x = rng.random(size=shape)
        y = (rng.random(size=shape) if y_continuous
             else rng.integers(0, 10, size=shape))
        method = stats.PermutationMethod(rng=rng)
        res = stats.chatterjeexi(x, y, method=method,
                                 y_continuous=y_continuous, axis=-1)
        ref = stats.chatterjeexi(x, y, y_continuous=y_continuous, axis=-1)
        np.testing.assert_allclose(res.statistic, ref.statistic, rtol=1e-15)
        np.testing.assert_allclose(res.pvalue, ref.pvalue, rtol=2e-2)

    def test_input_validation(self, xp):
        rng = np.random.default_rng(25932435798274926)
        x, y = rng.random(size=(2, 10))
        x, y = xp.asarray(x), xp.asarray(y)

        message = 'Array shapes are incompatible for broadcasting.|Incompatible shapes'
        with pytest.raises((ValueError, TypeError), match=message):
            stats.chatterjeexi(x, y[:-1])

        if not is_jax(xp):
            # jax misses out on some input validation from _axis_nan_policy decorator
            message = '...axis 10 is out of bounds for array...|out of range'
            with pytest.raises((ValueError, IndexError), match=message):
                stats.chatterjeexi(x, y, axis=10)

        message = '`y_continuous` must be boolean.'
        with pytest.raises(ValueError, match=message):
            stats.chatterjeexi(x, y, y_continuous='a herring')

        message = "`method` must be 'asymptotic' or"
        with pytest.raises(ValueError, match=message):
            stats.chatterjeexi(x, y, method='ekki ekii')

    @pytest.mark.skip_xp_backends('jax.numpy', reason='no SmallSampleWarning (lazy)')
    def test_special_cases(self, xp):
        message = 'One or more sample arguments is too small...'
        with pytest.warns(SmallSampleWarning, match=message):
            res = stats.chatterjeexi(xp.asarray([1]), xp.asarray([2]))

        assert xp.isnan(res.statistic)
        assert xp.isnan(res.pvalue)