Performa softmax on each segment. The first argument ``seglen`` stores the length of each segment. Its summation must be equal to the first dimension of the ``value`` tensor. Zero-length segments are allowed. Parameters ---------- seglen : Tensor Segment lengths.
(seglen, value)
| 60 | |
| 61 | |
| 62 | def segment_softmax(seglen, value): |
| 63 | """Performa softmax on each segment. |
| 64 | |
| 65 | The first argument ``seglen`` stores the length of each segment. Its |
| 66 | summation must be equal to the first dimension of the ``value`` tensor. |
| 67 | Zero-length segments are allowed. |
| 68 | |
| 69 | Parameters |
| 70 | ---------- |
| 71 | seglen : Tensor |
| 72 | Segment lengths. |
| 73 | value : Tensor |
| 74 | Value to aggregate. |
| 75 | |
| 76 | Returns |
| 77 | ------- |
| 78 | Tensor |
| 79 | Result tensor of the same shape as the ``value`` tensor. |
| 80 | |
| 81 | Examples |
| 82 | -------- |
| 83 | |
| 84 | >>> import dgl |
| 85 | >>> import torch as th |
| 86 | >>> val = th.ones(10, 3) |
| 87 | >>> seg = th.tensor([1, 0, 5, 4]) # 4 segments |
| 88 | >>> dgl.segment_softmax(seg, val) |
| 89 | tensor([[1.0000, 1.0000, 1.0000], |
| 90 | [0.2000, 0.2000, 0.2000], |
| 91 | [0.2000, 0.2000, 0.2000], |
| 92 | [0.2000, 0.2000, 0.2000], |
| 93 | [0.2000, 0.2000, 0.2000], |
| 94 | [0.2000, 0.2000, 0.2000], |
| 95 | [0.2500, 0.2500, 0.2500], |
| 96 | [0.2500, 0.2500, 0.2500], |
| 97 | [0.2500, 0.2500, 0.2500], |
| 98 | [0.2500, 0.2500, 0.2500]]) |
| 99 | """ |
| 100 | value_max = segment_reduce(seglen, value, reducer="max") |
| 101 | value = F.exp(value - F.repeat(value_max, seglen, dim=0)) |
| 102 | value_sum = segment_reduce(seglen, value, reducer="sum") |
| 103 | return value / F.repeat(value_sum, seglen, dim=0) |
| 104 | |
| 105 | |
| 106 | def segment_mm(a, b, seglen_a): |
nothing calls this directly
no test coverage detected