Skip to content

Compute _unique_2d on the device instead of through Python sets - #45

Open
MarcosAsh wants to merge 1 commit into
mainfrom
vectorize-unique-2d
Open

Compute _unique_2d on the device instead of through Python sets#45
MarcosAsh wants to merge 1 commit into
mainfrom
vectorize-unique-2d

Conversation

@MarcosAsh

@MarcosAsh MarcosAsh commented Aug 31, 2026

Copy link
Copy Markdown
Collaborator

Description

_unique_2d deduplicated rows with a Python set and built the inverse from a dict, calling .tolist() once per row each way. It runs inside ctc_decode's per timestep beam loop, so every step round tripped the whole path array through the host. This uses the device side lexsort unique already has in _unique_prepare instead.

It also fixes the inverse. The old code built unique_dict after padding, so when fill_value matched a genuinely unique row the padding entry won and the inverse pointed at a padding row. That is reachable from ctc_decode, which pads with -1 while an all -1 row is a real path.

This fixes the test_ctc_decode CI flake, which is the test behind four of the last five genuine CI failures. It segfaults in teardown after the test itself passes, and reproduces locally at about 2.5%. Running main and this branch interleaved, 600 runs each, gives 15 crashes on main and 0 here, one tailed p of 3e-5.

nn_test.py and numpy_test.py give 5750 passed and 0 failed, plus a differential test against the old implementation over 400 random inputs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant