Revert "fix: setitem in ak firsts with jax backend (#3455)" (#3460)

This reverts commit aeb9618.

Author: ianna

Diff for: src/awkward/operations/ak_firsts.py

@@ -10,7 +10,6 @@
     _named_axis_to_positional_axis,
     _remove_named_axis,
 )
-from awkward._nplikes.jax import Jax
 from awkward._nplikes.numpy_like import NumpyMetadata
 from awkward._nplikes.shape import unknown_length
 from awkward._regularize import regularize_axis
@@ -107,10 +106,7 @@ def action(layout, depth, backend, **kwargs):
                 stops = layout.stops.data
 
                 empties = starts == stops
-                if isinstance(layout.backend.nplike, Jax):
-                    index = index.at[empties].set(-1)
-                else:
-                    index[empties] = -1
+                index[empties] = -1
 
                 return ak.contents.IndexedOptionArray.simplified(
                     ak.index.Index64(index), layout._content