rllm.utils.seg_softmax

class rllm.utils.seg_softmax(data: Tensor, segment_ids: Tensor, num_segs: int)

Bases:

Compute the segment-wise softmax scores of elements in data.

Parameters:

• data (Tensor) – A tensor, typically two-dimensional.

• segment_ids (Tensor) – A one-dimensional tensor that indicates the segment assignment of each element in data.

• num_segs (int) – Total number of segments.

Returns:

Softmax scores with the same shape as data.

Return type:

Tensor