jax-ml/jax · error · ValueError

Sharded hypothesis runner does not support `data()` inside `

Error message

Sharded hypothesis runner does not support `data()` inside `@given`. All parameters must be drawn before the test body is called. Consider using `@composite` instead.

What it means

Raised by JAX's sharding-aware test harness (hypothesis_test_util) when a @given-based test class is configured for sharded execution but its @given includes hypothesis's data() strategy. Sharding needs all drawn parameters known before the test body runs, and data() defers draws into the body, so the combination is rejected with a suggestion to use @composite.

Source

Thrown at jax/_src/hypothesis_test_util.py:127


def _apply_sharding_to_tests(test_runner):

  shards_index_iter = itertools.cycle(range(_TEST_TOTAL_SHARDS))
  for name in dir(test_runner):
    if name.startswith("test"):
      test = getattr(test_runner, name)
      if detection.is_hypothesis_test(test):
        handle = test.hypothesis
        assert isinstance(handle, hp.core.HypothesisHandle)
        # `@given(..., data())` is not supported because:
        # - Sharding requires known values for all drawn parameters.
        # - The test body must be called given the sharding.
        # - Using `data()`, some or all parameters are not known until the test
        #   body is called.
        for val in handle._given_kwargs.values():
          if isinstance(val, hps_internal_core.DataStrategy):
            raise ValueError(
                "Sharded hypothesis runner does not support `data()` inside"
                " `@given`. All parameters must be drawn before the test body"
                " is called. Consider using `@composite` instead."
            )
        handle.inner_test = _shard_aware_hypothesis_inner_test(
            handle.inner_test
        )
      else:
        # If the tests are not hypothesis tests (or we are not sharding
        # hypothesis tests), we can just assign them to shards in a round-robin
        # fashion.
        if _TEST_TOTAL_SHARDS > 1:
          shard_index = next(shards_index_iter)
          setattr(test_runner, name, _shard_aware_test(test, shard_index))


class HypothesisShardedTestCase(jtu.JaxTestCase):
  """Runs Hypothesis tests in a sharded manner.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Replace data() with explicit strategies and @composite to generate dependent values up front
Draw all parameters in @given before the body runs
If sharding isn't needed for that test, run it outside the sharded runner/base class

Example fix

# before
@given(data=st.data())
def test_foo(self, data):
    x = data.draw(st.floats())
    ...
# after
@st.composite
def xs(draw):
    return draw(st.floats())
@given(x=xs())
def test_foo(self, x):
    ...

Defensive patterns

Strategy: validation

Validate before calling

# in test setup: reject data() before enabling sharding
from hypothesis.internal.core import DataStrategy  # illustrative
for k, strat in given_kwargs.items():
    if type(strat).__name__ == 'DataStrategy':
        raise ValueError('replace data() with @composite for sharded tests')

Prevention

Use @composite for dependent draws instead of data()
Keep @given draws static when sharding tests

When it happens

Trigger: Adding data() to the @given arguments of a test in a test class that inherits the sharded hypothesis runner (internal JAX test infra), e.g. @given(st.data(), x=st.floats()).

Common situations: Contributing tests to JAX (or projects reusing its test util) where sharding is enabled in CI; converting property tests to draw dependent values lazily via data().

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/0673f8993e25fb54. Report an issue: GitHub.