(gidx, logits, eids=ALL, norm_by="dst")
| 298 | |
| 299 | |
| 300 | def edge_softmax(gidx, logits, eids=ALL, norm_by="dst"): |
| 301 | @tf.custom_gradient |
| 302 | def _lambda(logits): |
| 303 | return edge_softmax_real(gidx, logits, eids, norm_by) |
| 304 | |
| 305 | return _lambda(logits) |
| 306 | |
| 307 | |
| 308 | def segment_reduce_real(op, x, offsets): |