Performs a sum-reduction of the features according the indices. output[indices[i]] += features[i] Args: features (layer) : 2D matrix with shape (E, F) indices (layer): 1D matrix with shape (E) dims (list of int): tuple of ints with the values (N, F)
(features, indices, dims, name=None)
| 16 | return lbann.Gather(features, indices, axis=0, name=name) |
| 17 | |
| 18 | def GraphReduce(features, indices, dims, name=None): |
| 19 | """Performs a sum-reduction of the features according the indices. |
| 20 | output[indices[i]] += features[i] |
| 21 | |
| 22 | Args: |
| 23 | features (layer) : 2D matrix with shape (E, F) |
| 24 | indices (layer): 1D matrix with shape (E) |
| 25 | dims (list of int): tuple of ints with the values (N, F) |
| 26 | returns: (layer) of shape (N, F) |
| 27 | """ |
| 28 | GraphReduce.count += 1 |
| 29 | if (name is None): |
| 30 | name = f"graph_reduce_{GraphReduce.count}" |
| 31 | return lbann.Scatter(features, indices, dims=dims, axis=0, name=name) |
| 32 | |
| 33 | GraphReduce.count = 0 |
| 34 | GraphExpand.count = 0 |