(op, x, offsets)
| 306 | |
| 307 | |
| 308 | def segment_reduce_real(op, x, offsets): |
| 309 | y, arg = _segment_reduce(op, x, offsets) |
| 310 | |
| 311 | def segment_reduce_backward(dy): |
| 312 | m = x.shape[0] |
| 313 | if op == "sum": |
| 314 | offsets_np = asnumpy(offsets[1:]) |
| 315 | indices_np = np.zeros((m + 1,), dtype=offsets_np.dtype) |
| 316 | np.add.at(indices_np, offsets_np, np.ones_like(offsets_np)) |
| 317 | indices_np = np.cumsum(indices_np, -1)[:-1] |
| 318 | indices = zerocopy_from_numpy(indices_np) |
| 319 | dx = tf.gather(dy, indices) |
| 320 | else: |
| 321 | dx = _bwd_segment_cmp(dy, arg, m) |
| 322 | return dx |
| 323 | |
| 324 | return y, segment_reduce_backward |
| 325 | |
| 326 | |
| 327 | def segment_reduce(op, x, offsets): |
no test coverage detected