<commit_before>from thinc.backends.jax_ops import lstm_weights_forward, backprop_lstm_weights
from thinc.backends.jax_ops import lstm_gates_forward, backprop_lstm_gates
import numpy.testing
import pytest

from hypothesis import given, settings
from ..strategies import ndarrays_of_shape

try:
    import jax

    has_jax = True
except ImportError:
    has_jax = False

MAX_EXAMPLES = 20

nL = 6
nB = 3
nO = 4
nI = 2
t = 3


def assert_arrays_equal(arrays1, arrays2):
    assert len(arrays1) == len(arrays2)
    shapes1 = [tuple(a.shape) for a in arrays1]
    shapes2 = [tuple(a.shape) for a in arrays2]
    assert shapes1 == shapes2
    for arr1, arr2 in zip(arrays1, arrays2):
        assert arr1.shape == arr2.shape
        numpy.testing.assert_allclose(arr1, arr2, rtol=0.001, atol=0.001)


# See thinc/backends/jax_ops for notation


@pytest.mark.skipif(not has_jax, reason="needs Jax")
@settings(max_examples=MAX_EXAMPLES, deadline=None)
@given(<commit_msg>Ignore JAX warning in test for cleaner output<commit_after>from thinc.backends.jax_ops import lstm_weights_forward, backprop_lstm_weights
from thinc.backends.jax_ops import lstm_gates_forward, backprop_lstm_gates
import numpy.testing
import pytest

from hypothesis import given, settings
from ..strategies import ndarrays_of_shape

try:
    import jax

    has_jax = True
except ImportError:
    has_jax = False

MAX_EXAMPLES = 20

nL = 6
nB = 3
nO = 4
nI = 2
t = 3


def assert_arrays_equal(arrays1, arrays2):
    assert len(arrays1) == len(arrays2)
    shapes1 = [tuple(a.shape) for a in arrays1]
    shapes2 = [tuple(a.shape) for a in arrays2]
    assert shapes1 == shapes2
    for arr1, arr2 in zip(arrays1, arrays2):
        assert arr1.shape == arr2.shape
        numpy.testing.assert_allclose(arr1, arr2, rtol=0.001, atol=0.001)


# See thinc/backends/jax_ops for notation


@pytest.mark.skipif(not has_jax, reason="needs Jax")
@pytest.mark.filterwarnings("ignore")
@settings(max_examples=MAX_EXAMPLES, deadline=None)
@given(
{'fraction_non_alphanumeric': 0.09077231695085256, 'fraction_numerical': 0.033099297893681046, 'mean_word_length': 4.377358490566038, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5155049', 'n_tokens_mistral': 765, 'n_tokens_neox': 729, 'n_words': 185}