mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #19540 from mattjj:remove-hypothesis-test-dependence

Browse Source 
PiperOrigin-RevId: 601908297
This commit is contained in:
jax authors 2024-01-26 17:28:26 -08:00
commit ccfe9c1ec2
1 changed files with 11 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
tests/pallas/splash_attention_kernel_test.py
View File

@ -21,8 +21,6 @@ from typing import Any, Callable, TypeVar
from absl.testing import absltest
from absl.testing import parameterized
import hypothesis as hp
import hypothesis.strategies as hps
import jax
from jax import random
from jax import tree_util
@ -33,6 +31,17 @@ import jax.numpy as jnp
import numpy as np
try:
import hypothesis as hp
import hypothesis.strategies as hps
CAN_USE_HYPOTHESIS = True
except (ModuleNotFoundError, ImportError):
CAN_USE_HYPOTHESIS = False
if not CAN_USE_HYPOTHESIS:
import sys; sys.exit(0)
jax.config.parse_flags_with_absl()
hp.settings.register_profile(