code: def unique_all(
      x: JaxArray,
      /,
      *,
      axis: Optional[int] = None,
      by_value: bool = True,
  ) -> Tuple[JaxArray, JaxArray, JaxArray, JaxArray]:
      Results = namedtuple(
          "Results",
          ["values", "indices", "inverse_indices", "counts"],
      )
  
      values, indices, inverse_indices, counts = jnp.unique(
          x,
          return_index=True,
          return_counts=True,
          return_inverse=True,
          axis=axis,
      )
  
      nan_count = jnp.sum(jnp.isnan(x)).item()
      if nan_count > 1:
          values = jnp.concatenate(
              (
                  values,
                  jnp.full(
                      fill_value=jnp.nan,
                      shape=(nan_count - 1,),
                      dtype=values.dtype,
                  ),
              ),
              axis=0,
          )
          counts = jnp.concatenate(
              (
                  counts[:-1],
                  jnp.full(
                      fill_value=1,
                      shape=(nan_count,),
                      dtype=counts.dtype,
                  ),
              ),
              axis=0,
          )
          nan_idx = jnp.where(jnp.isnan(x.flatten()))[0]
          indices = jnp.concatenate(
              (indices[:-1], nan_idx), axis=0
          ).astype(indices.dtype)
  
      if not by_value:
          sort_idx = jnp.argsort(indices)
          values = jnp.take(values, sort_idx, axis=axis)
          counts = jnp.take(counts, sort_idx)
          indices = jnp.take(indices, sort_idx)
          inv_sort_idx = ivy.current_backend().invert_permutation(
              sort_idx
          )
          inverse_indices = jnp.vectorize(
              lambda y: jnp.take(inv_sort_idx, y)
          )(inverse_indices)
  
      return Results(
          values.astype(x.dtype),
          indices,
          inverse_indices,
          counts,
      )
  
file_path: ivy/ivy/functional/backends/jax/set.py
imports: import jax.numpy
  from typing import Tuple, Optional
  from collections import namedtuple
  from ivy.functional.backends.jax import JaxArray
  import ivy
label_desc: Use image node_id 1 for calling a global function with example usage: unique_all() and returns: Results