Recommended way to do set based loss in flax?

[rdilip]

When we implement models such as DETR, we use an algorithm such as Hungarian matching to map the output of the model to the set of ground truths. What is the recommended way to do this in flax? One could implement Hungarian matching from scratch (I believe this is what google/Scenic does), but it would be ideal if there was a way to just use scipy.optimize.linear_sum_assignment. One can’t simply call the scipy function because the jax tracer arrays can’t be naively converted to numpy arrays.

[cgarciae]

Hey @rdilip, couple of thoughts:

• There is a jax.scipy.optimize module so it might be beneficial to also post this question in the JAX repo in case they can add tis.
• Optax hosts a couple of loss functions so maybe you can ask if they are willing to maintain this as well.
• If this loss function is implemented in Scenic it would be recommendable to just copy it directly into you own project.

[rdilip]

I'll post in the optax + jax repos -- ended up using scenic, as you suggested. Thanks!