How to share a bug on private code

[AdrienCorenflos]

Hi,

I am experiencing a genuine bug in JAX which can be summarised by "backward pass on scan segfaults, but runs fine if jit is globally disabled". However I'm struggling to extract a minimal reproducible example (as I have no clue what is going on in XLA), and I am not very willing to make my code completely public as I am writing it for a conf paper.

I'm however more than happy to give access to my repo to specific individuals who could have a look. Repro in there would be easy as it is essentially running a unittest.

Do you have policies of any kind in this regard?

Thanks,

Adrien

[jakevdp]

I would suggest doing your best to find a stripped-down reproduction that you can share. Even if you granted access to your code to one of the JAX team members, the chances of them being able to effectively debug and find the problem within a large & complicated code base are not high.

[AdrienCorenflos]

I understand. How do you recommend I go around isolating the bug to make a repro when I get a segmentation fault? The only stacktrace I get is this (I didn't see anything similar in the issues).

Fatal Python error: Segmentation fault

Current thread 0x00007ff02fb77740 (most recent call first):
File ".../venv/lib/python3.8/site-packages/jax/interpreters/xla.py", line 390 in _execute_compiled_primitive
File ".../venv/lib/python3.8/site-packages/jax/interpreters/xla.py", line 274 in apply_primitive
File ".../venv/lib/python3.8/site-packages/jax/core.py", line 610 in process_primitive
File ".../venv/lib/python3.8/site-packages/jax/core.py", line 265 in bind
File ".../venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1876 in scan_bind
File ".../venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1710 in _scan_transpose
File ".../venv/lib/python3.8/site-packages/jax/interpreters/ad.py", line 224 in backward_pass
File ".../venv/lib/python3.8/site-packages/jax/interpreters/ad.py", line 122 in unbound_vjp
File ".../venv/lib/python3.8/site-packages/jax/_src/api.py", line 2043 in _vjp_pullback_wrapper
File ".../venv/lib/python3.8/site-packages/jax/_src/api.py", line 878 in value_and_grad_f
File ".../venv/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162 in reraise_with_filtered_traceback
File ".../tests/test_standard_particle_filter.py", line 121 in _pf_with_correct_gradient_fwd
File ".../venv/lib/python3.8/site-packages/jax/linear_util.py", line 166 in call_wrapped
File ".../venv/lib/python3.8/site-packages/jax/interpreters/ad.py", line 362 in process_custom_vjp_call
File ".../venv/lib/python3.8/site-packages/jax/_src/custom_derivatives.py", line 571 in bind
File ".../venv/lib/python3.8/site-packages/jax/_src/custom_derivatives.py", line 491 in call
File ".../tests/test_standard_particle_filter.py", line 142 in loss_fun
File ".../venv/lib/python3.8/site-packages/jax/linear_util.py", line 166 in call_wrapped
File ".../venv/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 505 in trace_to_jaxpr
File ".../venv/lib/python3.8/site-packages/jax/interpreters/ad.py", line 102 in linearize
File ".../venv/lib/python3.8/site-packages/jax/interpreters/ad.py", line 115 in vjp
File ".../venv/lib/python3.8/site-packages/jax/_src/api.py", line 2136 in _vjp
File ".../venv/lib/python3.8/site-packages/jax/_src/api.py", line 2127 in vjp
File ".../venv/lib/python3.8/site-packages/jax/test_util.py", line 273 in check_vjp
File ".../venv/lib/python3.8/site-packages/jax/test_util.py", line 323 in _check_grads
File ".../venv/lib/python3.8/site-packages/jax/test_util.py", line 330 in check_grads
File ".../tests/test_standard_particle_filter.py", line 145 in test_bootstrap_gradients
File ".../venv/lib/python3.8/site-packages/_pytest/python.py", line 183 in pytest_pyfunc_call
File ".../venv/lib/python3.8/site-packages/pluggy/callers.py", line 187 in _multicall
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 84 in
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 93 in _hookexec
File ".../venv/lib/python3.8/site-packages/pluggy/hooks.py", line 286 in call
File ".../venv/lib/python3.8/site-packages/_pytest/python.py", line 1641 in runtest
File ".../venv/lib/python3.8/site-packages/_pytest/runner.py", line 162 in pytest_runtest_call
File ".../venv/lib/python3.8/site-packages/pluggy/callers.py", line 187 in _multicall
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 84 in
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 93 in _hookexec
File ".../venv/lib/python3.8/site-packages/pluggy/hooks.py", line 286 in call
File ".../venv/lib/python3.8/site-packages/_pytest/runner.py", line 255 in
File ".../venv/lib/python3.8/site-packages/_pytest/runner.py", line 311 in from_call
File ".../venv/lib/python3.8/site-packages/_pytest/runner.py", line 254 in call_runtest_hook
File ".../venv/lib/python3.8/site-packages/_pytest/runner.py", line 215 in call_and_report
File ".../venv/lib/python3.8/site-packages/_pytest/runner.py", line 126 in runtestprotocol
File ".../venv/lib/python3.8/site-packages/_pytest/runner.py", line 109 in pytest_runtest_protocol
File ".../venv/lib/python3.8/site-packages/pluggy/callers.py", line 187 in _multicall
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 84 in
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 93 in _hookexec
File ".../venv/lib/python3.8/site-packages/pluggy/hooks.py", line 286 in call
File ".../venv/lib/python3.8/site-packages/_pytest/main.py", line 348 in pytest_runtestloop
File ".../venv/lib/python3.8/site-packages/pluggy/callers.py", line 187 in _multicall
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 84 in
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 93 in _hookexec
File ".../venv/lib/python3.8/site-packages/pluggy/hooks.py", line 286 in call
File ".../venv/lib/python3.8/site-packages/_pytest/main.py", line 323 in _main
File ".../venv/lib/python3.8/site-packages/_pytest/main.py", line 269 in wrap_session
File ".../venv/lib/python3.8/site-packages/_pytest/main.py", line 316 in pytest_cmdline_main
File ".../venv/lib/python3.8/site-packages/pluggy/callers.py", line 187 in _multicall
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 84 in
File ".../venv/lib/python3.8/site-packages/pluggy/manager.py", line 93 in _hookexec
File ".../venv/lib/python3.8/site-packages/pluggy/hooks.py", line 286 in call
File ".../venv/lib/python3.8/site-packages/_pytest/config/init.py", line 162 in main
File "/home/adrien/pycharm-2020.3.3/plugins/python/helpers/pycharm/_jb_pytest_runner.py", line 51 in

Process finished with exit code 139 (interrupted by signal 11: SIGSEGV)

The most fishy bit seems to be

File ".../venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1876 in scan_bind
File ".../venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py", line 1710 in _scan_transpose

But I'm not sure where to go from there.

[jakevdp]

Could it possibly be an out-of-memory issue? I don't have any other insight on that, unfortunately. Let's leave this open and maybe someone else can be of more help.

[AdrienCorenflos]

Well, it works fine when I disable jitting, I am unittesting my simple case, with fairly small inputs, and my jaxpr seems reasonably short (1000 lines). If I give all the inputs to scan_p.bind in https://github.com/google/jax/blob/3827a8b26a60bc24d293343648759ccabd9742e5/jax/_src/lax/control_flow.py#L1710, is that enough for someone to hint how I can then make a python script repro?