code: def _segment_ids_to_np_indices(segment_ids):
      if len(segment_ids) < 1:
          return segment_ids
      shifted = np.roll(segment_ids, 1, axis=0)
      if JAX_MODE:
          shifted = shifted.at[0].set(-1)
      else:
          shifted[0] = -1
      return np.nonzero(segment_ids - shifted)[0]
  
imports: import collections
  import functools
  import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy.numpy_array import _reverse
  from tensorflow_probability.python.internal.backend.numpy.numpy_array import one_hot
  from tensorflow_probability.python.internal.backend.numpy.ops import _convert_to_tensor
label_desc: Use image node_id 12 for calling a global function with example usage: _segment_ids_to_np_indices(segment_ids) and returns: np, segment_ids