brainpy.math.segment_max#

brainpy.math.segment_max(data, segment_ids, num_segments=None, indices_are_sorted=False, unique_indices=False, bucket_size=None, mode=None)[source]#

segment_max operator for brainpy Array and Variable.

Parameters

  • data (Array) – An array with the values to be reduced.

  • segment_ids (Array) – An array with integer dtype that indicates the segments of data (along its leading axis) to be summed. Values can be repeated and need not be sorted.

  • num_segments (Optional, int) – An int with nonnegative value indicating the number of segments. The default is set to be the minimum number of segments that would support all indices in segment_ids, calculated as max(segment_ids) + 1. Since num_segments determines the size of the output, a static value must be provided to use segment_sum in a jit-compiled function.

  • indices_are_sorted (bool) – whether segment_ids is known to be sorted.

  • unique_indices (bool) – whether segment_ids is known to be free of duplicates.

  • bucket_size (int) – Size of bucket to group indices into. segment_sum is performed on each bucket separately to improve numerical stability of addition. Default None means no bucketing.

  • mode (lax.GatherScatterMode) – A jax.lax.GatherScatterMode value describing how out-of-bounds indices should be handled. By default, values outside of the range [0, num_segments) are dropped and do not contribute to the sum.

Returns

output – An array with shape (num_segments,) + data.shape[1:] representing the segment sums.

Return type

Array