phlash gpu throws assertion error

[percyfal]

I'm running phlash on three different datasets (three VCFs) from which I select a number of individuals for sub-population analyses, resulting in 9 datasets. 4 of these run to completion, but 5 fail, most often with the following assertion error:

Traceback (most recent call last):
  File "/crex/proj/naiss2023-6-254/private/users/perun/passer/.pixi/envs/phlash/lib/python3.12/site-packages/jax/_src/callback.py", line 100, in pure_callback_impl
    return tree_util.tree_map(np.asarray, callback(*args))
                                          ^^^^^^^^^^^^^^^
  File "/crex/proj/naiss2023-6-254/private/users/perun/passer/.pixi/envs/phlash/lib/python3.12/site-packages/jax/_src/callback.py", line 77, in __call__
    return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/crex/proj/naiss2023-6-254/private/users/perun/opt/phlash/src/phlash/gpu.py", line 409, in __call__
    self._compute_on_gpu(*args)
  File "/crex/proj/naiss2023-6-254/private/users/perun/opt/phlash/src/phlash/gpu.py", line 439, in _compute_on_gpu
    results[i] = gpu_kernel(pp, split_index, grad, barrier)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/crex/proj/naiss2023-6-254/private/users/perun/opt/phlash/src/phlash/gpu.py", line 293, in __call__
    assert np.all(ll < 0)
           ^^^^^^^^^^^^^^
AssertionError
2025-06-10 15:44:48,114; ERROR [jax._src.callback:pure_callback_impl]: jax.pure_callback failed

At first I thought this could be due to some anomaly when subsetting on individuals (e.g., lots of monomorphic sites, sites with only missing data), but I manually reran a population making sure only to include variant sites with zero missing data, and this error still occurs. Are there any requirements on the input data, or are any number of missing data fine?

For comparison, I also ran using phlash=1.0.3, without success.

[terhorst]

Thanks for another report! (I think you are among the first people outside my group to try and run this :-)

Any chance you can share one of the datasets you're getting the error on?

[percyfal]

I think we can share data, but finally that is up to @gbdnorden who owns the data to decide. If so, we can setup the file for download from our university server.

(And we know of several people here that use phlash, and I have done so successfully before, so it's odd I'm seeing these issues now 😅 )

[percyfal]

I have a dataset ready for download. I subset the original data and managed to reproduce a failure, albeit a different one. There is something weird with the data.

[LN-11]

I'm running phlash on three different datasets (three VCFs) from which I select a number of individuals for sub-population analyses, resulting in 9 datasets. 4 of these run to completion, but 5 fail, most often with the following assertion error:我正在三个不同的数据集（三个 VCF）上运行 phlash，我从中选择了一些个体进行亚种群分析，从而产生了 9 个数据集。其中 4 个运行完成，但 5 个失败，通常出现以下断言错误：

Traceback (most recent call last):
  File "/crex/proj/naiss2023-6-254/private/users/perun/passer/.pixi/envs/phlash/lib/python3.12/site-packages/jax/_src/callback.py", line 100, in pure_callback_impl
    return tree_util.tree_map(np.asarray, callback(*args))
                                          ^^^^^^^^^^^^^^^
  File "/crex/proj/naiss2023-6-254/private/users/perun/passer/.pixi/envs/phlash/lib/python3.12/site-packages/jax/_src/callback.py", line 77, in __call__
    return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/crex/proj/naiss2023-6-254/private/users/perun/opt/phlash/src/phlash/gpu.py", line 409, in __call__
    self._compute_on_gpu(*args)
  File "/crex/proj/naiss2023-6-254/private/users/perun/opt/phlash/src/phlash/gpu.py", line 439, in _compute_on_gpu
    results[i] = gpu_kernel(pp, split_index, grad, barrier)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/crex/proj/naiss2023-6-254/private/users/perun/opt/phlash/src/phlash/gpu.py", line 293, in __call__
    assert np.all(ll < 0)
           ^^^^^^^^^^^^^^
AssertionError
2025-06-10 15:44:48,114; ERROR [jax._src.callback:pure_callback_impl]: jax.pure_callback failed

At first I thought this could be due to some anomaly when subsetting on individuals (e.g., lots of monomorphic sites, sites with only missing data), but I manually reran a population making sure only to include variant sites with zero missing data, and this error still occurs. Are there any requirements on the input data, or are any number of missing data fine?起初我认为这可能是由于在个体上进行子集化时的一些异常（例如，许多单态位点、仅缺少数据的位点），但我手动重新运行了总体，确保仅包含缺失数据为零的变体位点，并且此错误仍然发生。对输入数据是否有任何要求，或者缺失数据的数量是否很好？

For comparison, I also ran using phlash=1.0.3, without success.为了进行比较，我还使用 phlash = 1.0.3 运行，但没有成功。

Hi, did you solve this problem?