From f34915c1e811bbb37e1d53e5de4f97ddd483d092 Mon Sep 17 00:00:00 2001
From: Paul O'Leary McCann <polm@dampfkraft.com>
Date: Sat, 10 Jul 2021 18:08:51 +0900
Subject: [PATCH] Use scatter_add to speed up span embed backprop

This was the slowest part of the code, and using scatter_add here
probably reduces the runtime by 50%.
---
 spacy/ml/models/coref.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/spacy/ml/models/coref.py b/spacy/ml/models/coref.py
index 719750ecb..66039564e 100644
--- a/spacy/ml/models/coref.py
+++ b/spacy/ml/models/coref.py
@@ -187,14 +187,10 @@ def span_embeddings_forward(
 
             out = model.ops.alloc2f(len(indoc), dim)
 
-            for ii, (start, end) in enumerate(dY.indices[offset:hi]):
-                # adjust indexes to align with doc
-                start -= tokoffset
-                end -= tokoffset
-
-                out[start] += starts[ii]
-                out[end] += ends[ii]
-                out[start:end] += spanvecs[ii]
+            idxs = dY.indices[offset:hi] - tokoffset
+            ops.scatter_add(out, idxs[:, 0], starts)
+            ops.scatter_add(out, idxs[:, 1], ends)
+            ops.scatter_add(out, idxs.T, spanvecs)
             oweights.append(out)
 
             offset = hi