可通过将Tensor转换为粗糙Tensor，其中tf.RaggedTensor.from_value_rowids使用segment_ids。然后沿着每个轴应用median。
例如，要获得上例的粗糙Tensor：

a = tf.RaggedTensor.from_value_rowids(
    values=x,
    value_rowids=[0, 1, 1, 1, 2])
# a.shape gives (3, None, 4) - Where None has 1,3,1 dimensions

然后，我们将中值应用于上述Tensor的每一行
1.每行Tensor的中值为：

def median(x):
    return tfp.stats.percentile(x[None,...], 50.0, interpolation='midpoint', axis=1)

1.段中值

def segment_median(data, segment_ids):
    ragged = tf.RaggedTensor.from_value_rowids(values=data, value_rowids=segment_ids)
    m = []
    for row in ragged:
        m.append(median(row))
    return tf.squeeze(tf.stack(m))

segment_median（x，tf.常量（[0，1，1，1，2]））返回

<tf.Tensor: shape=(3, 4), dtype=float32, numpy=
array([[0.1, 0.2, 0.4, 0.5],
       [0.1, 0.2, 0.2, 0.6],
       [0.1, 0.1, 0.6, 0.5]], dtype=float32)>