Mark xfail for leaked tracer tests

[fehiepsi]

This PR adds an entry in github action to track leaked tracer failing tests reported in #1981. I found the behaviors under JAX_CHECK_TRACER_LEAKS is a bit different from the normal behaviors so just want to perform a subset of tests.

Fixes #1981

[fehiepsi]

It seems that some issues were fixed so I'll just include the non-passing tests in the PR.

[tillahoffmann]

Looks great! Worth considering a custom marker for this or too much?

# conftest.py
def pytest_runtest_setup(item):
    marker = item.get_closest_marker("xfail_tracer_leak")
    if marker is not None and os.environ.get("JAX_CHECK_TRACER_LEAKS") == "1":
        item.add_marker(pytest.mark.xfail(reason="Expected tracer leak."))


# test files
@pytest.mark.xfail_tracer_leak
def test_...():
    ...

[fehiepsi]

Thanks for the suggestion! I prefer using pytest.mark it directly.

There are only 4 issues left (yay!), which I have created separate github issues for each of them.

• Tracking leaked in ASVGD #1982 It is clear where is the leak here
• Running mcmc within jit might cause tracer leak #2000 I'll need to think more how to address this issue. The caching logic is useful, but maybe it's not reliable under jit.
• Running svi with mutable states causes tracer leak #1999 Maybe not complicated, haven't thought about it yet.
• scan over discrete latent variables causes tracer leak #1998 I spent a bit of time on this before but couldn't track down the issue :(

While debugging, I found that there are many live arrays over the tests so I take this chance to remove all of them (and add a small check to conftest to assure that no live arrays will be created globally).

[tillahoffmann]

Looks great!