code: def segment_max(data, segment_ids, num_segments=None, **kwargs):
      data = convert_to_tensor(data)
      segment_ids = convert_to_tensor(segment_ids)
      num_repeats = torch.prod(
          torch.tensor(data.shape[1:], device=get_device())
      ).long()
      # To use `scatter_reduce` in torch, we need to replicate `segment_ids` into
      # the shape of `data`.
      segment_ids = (
          segment_ids.repeat_interleave(num_repeats)
          .view(*data.shape)
          .type(torch.int64)
      )
      num_segments = num_segments or len(torch.unique(segment_ids))
  
      # .scatter_reduce does not support -1 in the indices.
      # Add all out-of-bound indices value to an extra dimension after
      # num_segments, which is removed before returning the result.
  
      # Replacing the out-of-bound indices.
      segment_ids = torch.where(
          segment_ids >= 0, segment_ids, num_segments
      )
      segment_ids = torch.where(
          segment_ids < num_segments, segment_ids, num_segments
      )
  
      # Add one more dimension to the result shape with the "+1".
      shape = (num_segments + 1,) + tuple(data.shape[1:])
  
      result = torch.zeros(*shape, device=get_device()).scatter_reduce(
          0, segment_ids, data.float(), "amax"
      )
  
      # Removing the extra dimension.
      result = result[:-1, ...]
  
      return result.type(data.dtype)
  
imports: import math
  import torch
  from keras.backend import standardize_dtype
  from keras.backend.torch.core import convert_to_tensor
  from keras.backend.torch.core import get_device
  from keras.backend.torch.numpy import pad
label_desc: Use image node_id 2 for calling a global function with example usage: segment_max(data, segment_ids, num_segments) and returns: result