code: def segment_max(data, segment_ids, num_segments=None, sorted=False):
      if num_segments is None:
          num_segments = np.amax(segment_ids) + 1
  
      valid_indices = segment_ids >= 0  # Ignore segment_ids that are -1
      valid_data = data[valid_indices]
      valid_segment_ids = segment_ids[valid_indices]
  
      data_shape = list(valid_data.shape)
      data_shape[
          0
      ] = num_segments  # Replace first dimension (which corresponds to segments)
  
      if sorted:
          result = np.zeros(data_shape, dtype=valid_data.dtype)
          np.maximum.at(result, valid_segment_ids, valid_data)
      else:
          sort_indices = np.argsort(valid_segment_ids)
          sorted_segment_ids = valid_segment_ids[sort_indices]
          sorted_data = valid_data[sort_indices]
  
          result = np.zeros(data_shape, dtype=valid_data.dtype)
          np.maximum.at(result, sorted_segment_ids, sorted_data)
  
      return result
  
imports: import numpy
  from keras.backend import standardize_dtype
  from keras.backend.jax.math import fft
  from keras.backend.jax.math import fft2
  from keras.backend.numpy.core import convert_to_tensor
  from keras.utils.module_utils import scipy
label_desc: Use image node_id 2 for calling a global function with example usage: segment_max(data, segment_ids, num_segments, sorted) and returns: result