SegmentMean

Computes the mean along segments of a tensor.

Abstract Signature:

SegmentMean(data: Tensor, segment_ids: Tensor)

JAX (Core)

API: jax.ops.segment_mean
Strategy: Direct Mapping
Official Docs ↗

Keras

API: keras.ops.segment_max
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: tf.math.segment_mean
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: jax.ops.segment_mean
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: jax.ops.segment_mean
Strategy: Direct Mapping
Official Docs ↗