From cee3af580b29b4e3125a5596eb656799886fba45 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <jakevdp@google.com>
Date: Thu, 8 Apr 2021 12:11:14 -0700
Subject: [PATCH] Testing: avoid global fixture for doctests

---
 .github/workflows/ci-build.yaml | 2 ++
 conftest.py                     | 8 --------
 2 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/ci-build.yaml b/.github/workflows/ci-build.yaml
index 5dc81d82d..fcff0e3ba 100644
--- a/.github/workflows/ci-build.yaml
+++ b/.github/workflows/ci-build.yaml
@@ -141,6 +141,8 @@ jobs:
       run: |
         pip install -r docs/requirements.txt
     - name: Test documentation
+      env:
+        XLA_FLAGS: "--xla_force_host_platform_device_count=8"
       run: |
         pytest -n 1 docs
         pytest -n 1 --doctest-modules --ignore=jax/experimental/jax2tf jax
diff --git a/conftest.py b/conftest.py
index 423605ae2..d881cff15 100644
--- a/conftest.py
+++ b/conftest.py
@@ -15,7 +15,6 @@
 
 import jax
 import numpy
-import os
 import pytest
 
 
@@ -25,10 +24,3 @@ def add_imports(doctest_namespace):
   doctest_namespace["lax"] = jax.lax
   doctest_namespace["jnp"] = jax.numpy
   doctest_namespace["np"] = numpy
-
-
-@pytest.fixture(autouse=True)
-def spoof_devices(doctest_namespace):
-  # Set up runtime to mimic an 8-core machine
-  flags = os.environ.get('XLA_FLAGS', '')
-  os.environ['XLA_FLAGS'] = flags + " --xla_force_host_platform_device_count=8"