Repository: ivy-llc/ivy

Issue:
Using `jax.lax.while_loop` in the `jax` backend implementation for `ivy.nms`
The `ivy.nms` function is widely used for the transpilation of various models performing object detection.
As of today, `torch` and `tensorflow` have native framework functions for performing non-max suppression, namely `torchvision.ops.nms` and `tf.image.non_max_suppression`. Hence in the `tensorflow` backend implementation and in the `torchvision` sub-backend implementation we have used these 2 functions ([torch](https://github.com/unifyai/ivy/blob/91a0b62de5881dcf5e88949262780c252b030e28/ivy/functional/backends/torch/sub_backends/torchvision/layers.py#L43) and [tensorflow](https://github.com/unifyai/ivy/blob/91a0b62de5881dcf5e88949262780c252b030e28/ivy/functional/backends/tensorflow/layers.py#L837))
However, in the case of `jax` there isn't a native framework function that performs non-max suppression, so we currently have implemented the `jax` backend of `ivy.nms` using other `jax` functions as shown [here](https://github.com/unifyai/ivy/blob/91a0b62de5881dcf5e88949262780c252b030e28/ivy/functional/backends/jax/layers.py#L468). 
The goal here is to find a way to replace the pythonic `while` loop [here](https://github.com/unifyai/ivy/blob/91a0b62de5881dcf5e88949262780c252b030e28/ivy/functional/backends/jax/layers.py#L504) with `jax`'s own control flow using `jax.lax.while_loop` which should make the implementation more optimized and also help capture dynamism. This should be done ensuring that the tests for the function continue to pass!


Discussion:
@ivy-leaves assign me
Hey @mobley-trent 
You have picked up this issue 2 weeks ago but have not made any progress yet. Are you still working on it or could we assign this issue to someone else?

@NripeshN I have some work done on a local branch but I'm yet to push the code. Just making sure I have a working prototype
> [@NripeshN](https://github.com/NripeshN) I have some work done on a local branch but I'm yet to push the code. Just making sure I have a working prototype

Hey @mobley-trent 
Can you create a PR so we can track the progress? Feel free to make it a draft PR.

Resolving patch:
diff --git a/ivy/functional/backends/jax/layers.py b/ivy/functional/backends/jax/layers.py
--- a/ivy/functional/backends/jax/layers.py
+++ b/ivy/functional/backends/jax/layers.py
@@ -501,7 +501,8 @@ def nms(
         keep = jnp.zeros((size,), dtype=jnp.int64)
         keep_idx = 0
 
-        while jnp.unique(order).size > 1:
+        def body_fn(loop_vars):
+            keep, keep_idx, boxes, areas, order = loop_vars
             max_iou_idx = order[0]
             keep = keep.at[keep_idx].set(max_iou_idx)
             keep_idx += 1
@@ -522,6 +523,15 @@ def nms(
             boxes = boxes.at[forward].set(boxes[forward[::-1]])
             areas = areas.at[forward].set(areas[forward[::-1]])
 
+            return keep, keep_idx, boxes, areas, order
+
+        def cond_fn(loop_vars):
+            _, _, _, _, order = loop_vars
+            return jnp.min(order) != jnp.max(order)
+
+        init_vars = (keep, keep_idx, boxes, areas, order)
+        keep, keep_idx, boxes, _, _ = jlax.while_loop(cond_fn, body_fn, init_vars)
+
         ret = jnp.array(keep[:keep_idx], dtype=jnp.int64)
 
     if len(ret) > 1 and scores is not None:


Test patch: