"""
Compare tagging speed for LSTM, using dummy data.

Results on CPU laptop:

PyTorchLSTM.v0:
Predicted 39017 4.033804399892688 Ys[0] 0.05000001 5.551115e-17

LSTM (NumpyOps):
Predicted 39018 13.174870599992573 Ys[0] 0.05000001 5.551115e-17

So PyTorch is 3x faster currently.
"""
import typer
from timeit import default_timer as timer
import numpy.random
from timeit import default_timer as timer
from typing import List, cast
from thinc.types import Array2d, Padded
import thinc.api
from thinc.api import Model, Config, registry, Embed, LSTM, Softmax
from thinc.api import chain, list2padded, with_array, to_categorical
from thinc.api import minibatch
from thinc.backends import jax_jit
import jax.tree_util

CONFIG = """
[data]
n_samples = 20000
n_tags = 20
n_vocab = 10000
length_mean = 40
length_variance = 5

[common]
width = 300

[model]
@layers = "LSTMTagger.v0"

[model.embed]
@layers = "Embed.v0"
nO = ${common:width}
nV = ${data:n_vocab}

[model.encode]
@layers = "LSTM.v0"
nO = ${common:width}
nI = ${common:width}
depth = 2

[model.predict]
@layers = "Softmax.v0"
nO = ${data:n_tags}
"""

@registry.layers("LSTMTagger.v0")
def build_tagger(
        embed: Model[Array2d, Array2d],
        encode: Model[Padded, Padded],
        predict: Model[Array2d, Array2d]
) -> Model[List[Array2d], Padded]:
    model = chain(
        list2padded(),
        with_array(embed),
        encode,
        #with_array(predict),
    )
    model.set_ref("embed", embed)
    model.set_ref("encode", encode)
    model.set_ref("predict", model.layers[-1])
    return model


def get_dummy_data(n_samples, n_tags, n_vocab, length_mean, length_variance):
    Xs = []
    Ys = []
    for _ in range(n_samples):
        length = numpy.random.normal(size=1, scale=length_variance) + length_mean
        shape = (max(1, int(length)),)
        X = numpy.random.uniform(0, n_vocab-1, shape)
        Y = numpy.random.uniform(0, n_tags-1, shape)
        assert X.size, length
        assert Y.size, length
        Xs.append(X.reshape((-1, 1)).astype("i"))
        Ys.append(to_categorical(Y.astype("i")))
    return Xs, Ys

jax.tree_util.register_pytree_node(
    Padded,
    lambda pad: ((pad.data, pad.size_at_t, pad.lengths, pad.indices), None),
    lambda info, values: Padded(*values)
)

def run_forward(model, Xs):
    Ys = []
    for batch in Xs:
        Ys.append(model.predict(batch))
    for Y in Ys:
        if hasattr(Y.data, "block_until_ready"):
            Y.data.block_until_ready()
    return Ys

def set_backend(name, gpu_id):
    global CONFIG
    if name == "jax":
        thinc.api.set_current_ops(thinc.api.JaxOps())
        CONFIG = CONFIG.replace("PyTorch", "")
    else:
        if gpu_id == -1:
            thinc.api.set_current_ops(thinc.api.NumpyOps())
        else:
            thinc.api.set_current_ops(thinc.api.CupyOps())
        CONFIG = CONFIG.replace("LSTM.v0", "PyTorchLSTM.v0")


def main(jax: bool=False, pytorch: bool=False, gpu_id: int=-1):
    global CONFIG
    thinc.api.fix_random_seed(0)
    if gpu_id >= 0:
        thinc.api.require_gpu(gpu_id)
    backends = {"jax": jax, "pytorch": pytorch}
    for name, use_backend in backends.items():
        if not use_backend:
            print(f"Skipping {name}")
            continue
        set_backend(name, gpu_id)
        C = registry.make_from_config(Config().from_str(CONFIG))
        model = C["model"]
        X, Y = get_dummy_data(**C["data"])
        print("Begin init", len(X))
        model.initialize(X=X[:5])
        print("Pre-batch")
        n_words = sum(len(x) for x in X)
        X = [model.layers[0].predict(batch) for batch in minibatch(X, size=128)]
        model.layers.pop(0)
        print("Start")
        start_time = timer()
        Ys = run_forward(model, X)
        end_time = timer()
        print(name, n_words, end_time-start_time)


if __name__ == "__main__":
    typer.run(main)

@@ -737,17 +737,18 @@
 mples = 
-2
+10
 0000%0An_t
@@ -807,17 +807,17 @@
 iance = 
-5
+1
 %0A%0A%5Bcommo
@@ -1030,17 +1030,17 @@
 depth = 
-2
+1
 %0A%0A%5Bmodel
@@ -1702,16 +1702,17 @@
         
+#
 length =
@@ -1777,16 +1777,45 @@
 th_mean%0A
+        length = length_mean%0A
         
@@ -2327,31 +2327,34 @@
 l, Xs):%0A    
-Ys = %5B%5D
+total = 0.
 %0A    for bat
@@ -2372,25 +2372,19 @@
        Y
-s.append(
+ = 
 model.pr
@@ -2399,112 +2399,36 @@
 tch)
-)%0A    for Y in Ys:%0A        if hasattr(Y.data, %22block_until_ready%22):%0A            Y.data.block_until_ready
+%0A        total += Y.data.sum
 ()%0A 
@@ -2437,18 +2437,28 @@
  return 
-Ys
+float(total)
 %0A%0Adef se
@@ -3023,24 +3023,57 @@
 gpu(gpu_id)%0A
+        print(%22Set GPU%22, gpu_id)%0A
     backends
@@ -3411,16 +3411,140 @@
 data%22%5D)%0A
+        print(%22Copy to device%22)%0A        X = %5Bmodel.ops.asarray(x) for x in X%5D%0A        Y = %5Bmodel.ops.asarray(y) for y in Y%5D%0A
         
@@ -3753,10 +3753,9 @@
 ze=1
-28
+6
 )%5D%0A 
@@ -3841,18 +3841,21 @@
         
-Ys
+total
  = run_f

Update benchmark