import pytest
from numpy.testing import assert_allclose
import thinc.api
from thinc.api import has_jax
from thinc.backends import jax_jit


def get_batch(ops, nB, nO, nI):
    X = ops.xp.zeros((nB, nI), dtype="f")
    X += ops.xp.random.uniform(-1, 1, X.shape)
    Y = ops.xp.zeros((nB, nO), dtype="f")
    Y += ops.xp.random.uniform(-1, 1, Y.shape)
    return X, Y


def make_linear(nO, nI):
    model = thinc.api.Linear(nO, nI)
    model.set_attr("registry_name", "Linear.v0")
    return model


@jax_jit()
def accepts_thinc_model(model):
    return model.get_param("W").sum()


@pytest.mark.skipif(not has_jax, reason="needs Jax")
def test_jax_jit_function_accepts_model():
    thinc.api.set_current_ops(thinc.api.JaxOps())
    model = make_linear(2, 2)
    sum_W = accepts_thinc_model(model)
    assert_allclose(float(sum_W), float(model.get_param("W").sum()), atol=1e-4)


@pytest.mark.skipif(not has_jax, reason="needs Jax")
def test_jax_jit_linear_forward(nB=8, nI=4, nO=3):
    thinc.api.set_current_ops(thinc.api.JaxOps())
    model = make_linear(nO=nO, nI=nI)
    X, Y = get_batch(model.ops, nB=nB, nO=nO, nI=nI)
    Yh = model.predict(X)
    model._func = jax_jit()(model._func)
    Yh_jit = model.predict(X)
    assert_allclose(Yh, Yh_jit)


@pytest.mark.skipif(not has_jax, reason="needs Jax")
def test_jax_jit_static_arg_linear_forward(nB=8, nI=4, nO=3):
    thinc.api.set_current_ops(thinc.api.JaxOps())
    model = make_linear(nO=nO, nI=nI)
    X, Y = get_batch(model.ops, nB=nB, nO=nO, nI=nI)
    Yh = model.predict(X)
    model._func = jax_jit(0)(model._func)
    Yh_jit = model.predict(X)
    assert_allclose(Yh, Yh_jit)

@@ -7,16 +7,34 @@
  pytest%0A
+import contextlib%0A
 from num
@@ -150,16 +150,190 @@
 ax_jit%0A%0A
+@contextlib.contextmanager%0Adef with_current_ops(ops):%0A    prev = thinc.api.get_current_ops()%0A    thinc.api.set_current_ops(ops)%0A    yield%0A    thinc.api.set_current_ops(prev)%0A
 %0Adef get
@@ -862,37 +862,33 @@
 odel():%0A    
-thinc.api.set
+with with
 _current_ops
@@ -899,33 +899,38 @@
 nc.api.JaxOps())
-%0A
+:%0A    
     model = make
@@ -947,16 +947,20 @@
  2)%0A    
+  
+  
 sum_W = 
@@ -986,16 +986,20 @@
 (model)%0A
+    
     asse
@@ -1176,37 +1176,33 @@
  nO=3):%0A    
-thinc.api.set
+with with
 _current_ops
@@ -1213,37 +1213,42 @@
 nc.api.JaxOps())
-%0A
+:%0A   
     
+ 
 model = make_lin
@@ -1257,32 +1257,36 @@
 r(nO=nO, nI=nI)%0A
+    
     X, Y = get_b
@@ -1314,32 +1314,36 @@
 , nO=nO, nI=nI)%0A
+    
     Yh = model.p
@@ -1344,32 +1344,36 @@
 odel.predict(X)%0A
+    
     model._func 
@@ -1389,32 +1389,36 @@
 ()(model._func)%0A
+    
     Yh_jit = mod
@@ -1423,32 +1423,36 @@
 odel.predict(X)%0A
+    
     assert_allcl
@@ -1580,24 +1580,75 @@
 I=4, nO=3):%0A
+    with with_current_ops(thinc.api.JaxOps()):%0A    
     thinc.ap
@@ -1685,24 +1685,28 @@
 xOps())%0A    
+   
+ 
 model = make
@@ -1727,24 +1727,28 @@
  nI=nI)%0A    
+   
+ 
 X, Y = get_b
@@ -1780,24 +1780,28 @@
 =nO, nI=nI)%0A
+    
     Yh = mod
@@ -1810,24 +1810,28 @@
 .predict(X)%0A
+    
     model._f
@@ -1856,24 +1856,28 @@
 odel._func)%0A
+    
     Yh_jit =
@@ -1890,24 +1890,28 @@
 .predict(X)%0A
+    
     assert_a

Fix jax model test