Sphinx 7.2 breaks Matplotlib's directive with Jax's arrays

[jeertmans]

Describe the bug

Hello,

When using the usual make html, I faced the following issue (minimized):

...
TypeError: concatenate requires ndarray or scalar arguments, got <class 'jaxlib.xla_extension.ArrayImpl'> at position 0.
/export/home/eertmans/repositories/DiffeRT/DiffeRT2d/differt2d/scene.py:docstring of differt2d.scene.Scene.square_scene:20: WARNING: Exception occurred in plotting differt2d-scene-3
 from /export/home/eertmans/repositories/DiffeRT/DiffeRT2d/docs/source/reference/differt2d.scene.rst:
Traceback (most recent call last):
  File "/home/eertmans/.cache/pypoetry/virtualenvs/differt2d-4IUpg1Dz-py3.10/lib/python3.10/site-packages/matplotlib/sphinxext/plot_directive.py", line 481, in _run_code
    exec(code, ns)
  File "<string>", line 6, in <module>
  File "/export/home/eertmans/repositories/DiffeRT/DiffeRT2d/differt2d/scene.py", line 602, in plot
    [
  File "/export/home/eertmans/repositories/DiffeRT/DiffeRT2d/differt2d/scene.py", line 603, in <listcomp>
    emitter.plot(ax, *emitters_args, *args, **emitters_kwargs, **kwargs)
  File "/export/home/eertmans/repositories/DiffeRT/DiffeRT2d/differt2d/geometry.py", line 254, in plot
    [self.point[0]],
TypeError: 'ArrayImpl' object is not subscriptable
/export/home/eertmans/repositories/DiffeRT/DiffeRT2d/differt2d/scene.py:docstring of differt2d.scene.Scene.square_scene_with_obstacle:24: WARNING: Exception occurred in plotting differt2d-scene-4
 from /export/home/eertmans/repositories/DiffeRT/DiffeRT2d/docs/source/reference/differt2d.scene.rst:
Traceback (most recent call last):
  File "/home/eertmans/.cache/pypoetry/virtualenvs/differt2d-4IUpg1Dz-py3.10/lib/python3.10/site-packages/matplotlib/sphinxext/plot_directive.py", line 481, in _run_code
    exec(code, ns)
  File "<string>", line 6, in <module>
  File "/export/home/eertmans/repositories/DiffeRT/DiffeRT2d/differt2d/scene.py", line 602, in plot
    [
  File "/export/home/eertmans/repositories/DiffeRT/DiffeRT2d/differt2d/scene.py", line 603, in <listcomp>
    emitter.plot(ax, *emitters_args, *args, **emitters_kwargs, **kwargs)
  File "/export/home/eertmans/repositories/DiffeRT/DiffeRT2d/differt2d/geometry.py", line 254, in plot
    [self.point[0]],
TypeError: 'ArrayImpl' object is not subscriptable
looking for now-outdated files... none found
pickling environment... done
checking consistency... done
preparing documents... done
                  copying downloadable files... [  8%] ../build/plot_directive/reference/differt2d-geo
copying downloadable files... [100%] ../build/plot_directive/reference/differt2d-scene-4.py
copying static files... done
copying extra files... done
done
writing output... [100%] reference/index
generating indices... genindex py-modindex done
highlighting module code... [100%] differt2d.utils
writing additional pages... search done
dumping search index in English (code: en)... done
dumping object inventory... done
build succeeded, 13 warnings.

The HTML pages are in build/html.

As you can see, they say that ArrayImpl is not indexable. But self.point (and others) should not be a ArrayImpl, but rather some concrete array value. When running pytest and other commands, I get the expected output. When building with Sphinx<7.2, it works as expected. I don't know what triggers this ArrayImpl type to be used, and this may be a side effect of using Jax.

This issue is faced when combining Sphinx, jax and matplotlib's Sphinx directive altogether. The error was quite hard to pinpoint, and it is still quite vague to me, even after taking a look at the 7.2 release notes.

I think the issue comes from the Sphinx side, because changing the version requirements to Sphinx<7.2 fixed it all. Using pip freeze, there was no other relevant package that could cause this issue (to my understanding). This is why I post the issue in this repository, but please tell me if I just rather put it elsewhere.

How to Reproduce

This is not a real MWE, but the issue happened on my repo (https://github.com/jeertmans/DiffeRT2d), which you can see the different between this commit (not working) to this commit (working).

Build logs are available here for the failing version:
logs_281.zip

Environment Information

Platform:              linux; (Linux-6.2.0-26-generic-x86_64-with-glibc2.35)
Python version:        3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0])
Python implementation: CPython
Sphinx version:        7.2.0
Docutils version:      0.20.1
Jinja2 version:        3.1.2
Pygments version:      2.16.1

Sphinx extensions

extensions = [
      # Built-in
      "sphinx.ext.autodoc",
      "sphinx.ext.extlinks",
      "sphinx.ext.intersphinx",
      "sphinx.ext.mathjax",
      "sphinx.ext.viewcode",
      # Additional
      "myst_parser",
      "matplotlib.sphinxext.plot_directive",  # Relevant one
      "sphinxext.opengraph",
      "sphinx_copybutton",
      "sphinx_autodoc_typehints",
]

Additional context

No response

[AA-Turner]

Hi Jérome (@jeertmans),

Please would you be able to paste the full log mentioned in the failure message, and ideally also the untruncated -T output?

Additionally, please may you try to reduce the set of extensions as much as possible -- certainly removing myst_parser, and ideally as much as possible. If this is an autodoc failure, ideally just a single .. auto<blah>:: directive in an index.rst file would be very helpful, otherwise knowing which directive caused the failure etc, please.

Thanks,
Adam

[jeertmans]

Ok so I managed to have a MWE. The error is not the same, but it also occurs when upgrading to Sphinx 7.2.

pyproject.toml

[tool.poetry]
name = "sphinx-11652-mwe"
version = "0.1.0"
description = ""
authors = ["Jérome Eertmans <jeertmans@icloud.com>"]
readme = "README.md"
packages = [{include = "sphinx_11652_mwe"}]

[tool.poetry.dependencies]
python = "^3.10"
sphinx = "<7.2" # Works
# sphinx = ">=7.2" # Fails
jax = {extras = ["cpu"], version = "^0.4.14"}
matplotlib = "^3.7.2"
chex = "^0.1.82"

[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"

sphinx_11652_mwe/__init__.py

from __future__ import annotations

import jax.numpy as jnp
from chex import dataclass
from jax import Array


@dataclass
class Foo:
    """
    .. plot::

        import matplotlib.pyplot as plt
        from sphinx_11652_mwe import Foo

        n = 5

        f = Foo.from_size(n)

        plt.plot(f.x[n:], f.x[:n])
        plt.show()
    """
    x: Array

    @classmethod
    def from_size(cls, n: int) -> Foo:
        return cls(x=jnp.arange(2 * n))

source/conf.py

project = 'test'
copyright = '2023, NA'
author = 'NA'

extensions = [
    'sphinx.ext.autodoc',
    'matplotlib.sphinxext.plot_directive'
]

templates_path = ['_templates']
exclude_patterns = []

html_theme = 'alabaster'
html_static_path = ['_static']

source/index.rst

Welcome to test's documentation!
================================

.. automodule:: sphinx_11652_mwe
   :members:

Output from failing build:

Full traceback

WARNING: autodoc: failed to import module 'sphinx_11652_mwe'; the following exception was raised:
Traceback (most recent call last):
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 64, in import_module
    return importlib.import_module(modname)
  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/export/home/eertmans/Desktop/sphinx-11652-mwe/sphinx_11652_mwe/__init__.py", line 3, in <module>
    import jax.numpy as jnp
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/__init__.py", line 169, in <module>
    from jax import scipy as scipy
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/scipy/__init__.py", line 24, in <module>
    from jax.scipy import signal as signal
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/scipy/signal.py", line 18, in <module>
    from jax._src.scipy.signal import (
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/_src/scipy/signal.py", line 23, in <module>
    import scipy.signal as osp_signal
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/signal/__init__.py", line 323, in <module>
    from ._filter_design import *
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/signal/_filter_design.py", line 16, in <module>
    from scipy import special, optimize, fft as sp_fft
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/__init__.py", line 211, in __getattr__
    return _importlib.import_module(f'scipy.{name}')
  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/optimize/__init__.py", line 424, in <module>
    from ._direct_py import direct
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/optimize/_direct_py.py", line 13, in <module>
    from _typeshed import NoneType
ModuleNotFoundError: No module named '_typeshed'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 88, in import_object
    module = import_module(modname, warningiserror=warningiserror)
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 68, in import_module
    raise ImportError(exc, traceback.format_exc()) from exc
ImportError: (ModuleNotFoundError("No module named '_typeshed'"), 'Traceback (most recent call last):\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 64, in import_module\n    return importlib.import_module(modname)\n  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module\n    return _bootstrap._gcd_import(name[level:], package, level)\n  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import\n  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load\n  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked\n  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked\n  File "<frozen importlib._bootstrap_external>", line 883, in exec_module\n  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed\n  File "/export/home/eertmans/Desktop/sphinx-11652-mwe/sphinx_11652_mwe/__init__.py", line 3, in <module>\n    import jax.numpy as jnp\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/__init__.py", line 169, in <module>\n    from jax import scipy as scipy\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/scipy/__init__.py", line 24, in <module>\n    from jax.scipy import signal as signal\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/scipy/signal.py", line 18, in <module>\n    from jax._src.scipy.signal import (\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/_src/scipy/signal.py", line 23, in <module>\n    import scipy.signal as osp_signal\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/signal/__init__.py", line 323, in <module>\n    from ._filter_design import *\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/signal/_filter_design.py", line 16, in <module>\n    from scipy import special, optimize, fft as sp_fft\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/__init__.py", line 211, in __getattr__\n    return _importlib.import_module(f\'scipy.{name}\')\n  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module\n    return _bootstrap._gcd_import(name[level:], package, level)\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/optimize/__init__.py", line 424, in <module>\n    from ._direct_py import direct\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/optimize/_direct_py.py", line 13, in <module>\n    from _typeshed import NoneType\nModuleNotFoundError: No module named \'_typeshed\'\n')

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 64, in import_module
    return importlib.import_module(modname)
  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/export/home/eertmans/Desktop/sphinx-11652-mwe/sphinx_11652_mwe/__init__.py", line 4, in <module>�[91mWARNING: autodoc: failed to import module 'sphinx_11652_mwe'; the following exception was raised:
Traceback (most recent call last):
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 64, in import_module
    return importlib.import_module(modname)
  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/export/home/eertmans/Desktop/sphinx-11652-mwe/sphinx_11652_mwe/__init__.py", line 3, in <module>
    import jax.numpy as jnp
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/__init__.py", line 169, in <module>
    from jax import scipy as scipy
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/scipy/__init__.py", line 24, in <module>
    from jax.scipy import signal as signal
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/scipy/signal.py", line 18, in <module>
    from jax._src.scipy.signal import (
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/_src/scipy/signal.py", line 23, in <module>
    import scipy.signal as osp_signal
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/signal/__init__.py", line 323, in <module>
    from ._filter_design import *
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/signal/_filter_design.py", line 16, in <module>
    from scipy import special, optimize, fft as sp_fft
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/__init__.py", line 211, in __getattr__
    return _importlib.import_module(f'scipy.{name}')
  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/optimize/__init__.py", line 424, in <module>
    from ._direct_py import direct
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/optimize/_direct_py.py", line 13, in <module>
    from _typeshed import NoneType
ModuleNotFoundError: No module named '_typeshed'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 88, in import_object
    module = import_module(modname, warningiserror=warningiserror)
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 68, in import_module
    raise ImportError(exc, traceback.format_exc()) from exc
ImportError: (ModuleNotFoundError("No module named '_typeshed'"), 'Traceback (most recent call last):\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 64, in import_module\n    return importlib.import_module(modname)\n  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module\n    return _bootstrap._gcd_import(name[level:], package, level)\n  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import\n  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load\n  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked\n  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked\n  File "<frozen importlib._bootstrap_external>", line 883, in exec_module\n  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed\n  File "/export/home/eertmans/Desktop/sphinx-11652-mwe/sphinx_11652_mwe/__init__.py", line 3, in <module>\n    import jax.numpy as jnp\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/__init__.py", line 169, in <module>\n    from jax import scipy as scipy\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/scipy/__init__.py", line 24, in <module>\n    from jax.scipy import signal as signal\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/scipy/signal.py", line 18, in <module>\n    from jax._src.scipy.signal import (\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/_src/scipy/signal.py", line 23, in <module>\n    import scipy.signal as osp_signal\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/signal/__init__.py", line 323, in <module>\n    from ._filter_design import *\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/signal/_filter_design.py", line 16, in <module>\n    from scipy import special, optimize, fft as sp_fft\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/__init__.py", line 211, in __getattr__\n    return _importlib.import_module(f\'scipy.{name}\')\n  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module\n    return _bootstrap._gcd_import(name[level:], package, level)\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/optimize/__init__.py", line 424, in <module>\n    from ._direct_py import direct\n  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/scipy/optimize/_direct_py.py", line 13, in <module>\n    from _typeshed import NoneType\nModuleNotFoundError: No module named \'_typeshed\'\n')

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/sphinx/ext/autodoc/importer.py", line 64, in import_module
    return importlib.import_module(modname)
  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/export/home/eertmans/Desktop/sphinx-11652-mwe/sphinx_11652_mwe/__init__.py", line 4, in <module>
    from chex import dataclass
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/chex/__init__.py", line 17, in <module>
    from chex._src.asserts import assert_axis_dimension
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/chex/_src/asserts.py", line 27, in <module>
    from chex._src import asserts_internal as _ai
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/chex/_src/asserts_internal.py", line 34, in <module>
    from chex._src import pytypes
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/chex/_src/pytypes.py", line 53, in <module>
    Shape = jax.core.Shape
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/_src/deprecations.py", line 53, in getattr
    raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax' has no attribute 'core'
�[39;49;00m
    from chex import dataclass
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/chex/__init__.py", line 17, in <module>
    from chex._src.asserts import assert_axis_dimension
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/chex/_src/asserts.py", line 27, in <module>
    from chex._src import asserts_internal as _ai
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/chex/_src/asserts_internal.py", line 34, in <module>
    from chex._src import pytypes
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/chex/_src/pytypes.py", line 53, in <module>
    Shape = jax.core.Shape
  File "/home/eertmans/.cache/pypoetry/virtualenvs/sphinx-11652-mwe-IiPyajJx-py3.10/lib/python3.10/site-packages/jax/_src/deprecations.py", line 53, in getattr
    raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax' has no attribute 'core'

I have attached the project files if you prefer.
sphinx-11652-mwe.zip

[picnixz]

Ah it's maybe because of some TYPE_CHECKING guard I think. It appears to me that the original exception is the fact that _typeshed cannot be found and this appears to come from scipy.

[jeertmans]

That's what I thought, too, but what change in 7.2 could have cause that?

[AA-Turner]

@jeertmans -- I'm not able to reproduce the error you're seeing. I get:

/mnt/s/Development/sphinx-issue-11652/sphinx_11652_mwe.py:docstring of sphinx_11652_mwe.Foo:1: WARNING: Exception occurred in plotting index-1
 from /mnt/s/Development/sphinx-issue-11652/index.rst:
Traceback (most recent call last):
  File "/mnt/s/Development/sphinx-issue-11652/.venv/lib/python3.11/site-packages/matplotlib/sphinxext/plot_directive.py", line 481, in _run_code
    exec(code, ns)
  File "<string>", line 8, in <module>
TypeError: 'ArrayImpl' object is not subscriptable

with:

Heading
=======

.. automodule:: sphinx_11652_mwe
   :members:

import os, sys
sys.path.insert(0, os.path.abspath('.'))

extensions = [
    'sphinx.ext.autodoc',
    'matplotlib.sphinxext.plot_directive'
]

exclude_patterns = [
    '.venv',
]

html_theme = 'alabaster'

from __future__ import annotations

import jax.numpy as jnp
from chex import dataclass
from jax import Array

@dataclass
class Foo:
    """
    .. plot::

        import matplotlib.pyplot as plt
        from sphinx_11652_mwe import Foo

        f = Foo.from_size(5)

        plt.plot(f.x[5:], f.x[:5])
        plt.show()
    """
    x: Array

    @classmethod
    def from_size(cls, n: int) -> Foo:
        return cls(x=jnp.arange(2 * n))

Are you able to check without Poetry please? I ran: python3.11 -m venv .venv and then pip install sphinx jax[cpu] chex matplotlib in the activated venv (through WSL). I then ran Sphinx with sphinx-build -M html . _build -aET.

A

[jeertmans]

I get the exact same error message as you when I do not use poetry. But what is also weird is that when I pip freeze and compare both environments (Poetry and -m venv), I get the exact same output.

Anyway, having ArrayImpl instead of a concrete array is still an issue...

[picnixz]

Not entirely sure, but I'm wondering how matplotlib actually executes the code. Maybe it is an issue with what the local namespace consists of? Also, could you like print the type of f.x and check what happens? (I never used the plot directive) and check the output of jnp.arange(...) in the classmethod itself.

Finally, is the chex.dataclass possibly responsible for this? can you use a real dataclass instead and real functions to create your dataclass (still with this Array annotation)?

[jeertmans]

I guess so..

Printing f.x (by raising a ValueError, because Matplotlib captures the output otherwise):

/home/jeertmans/Downloads/sphinx-11652-mwe/sphinx_11652_mwe/__init__.py:docstring of sphinx_11652_mwe.Foo:1: WARNING: Exception occurred in plotting index-1
 from /home/jeertmans/Downloads/sphinx-11652-mwe/source/index.rst:
Traceback (most recent call last):
  File "/home/jeertmans/Downloads/sphinx-11652-mwe/venv/lib/python3.10/site-packages/matplotlib/sphinxext/plot_directive.py", line 481, in _run_code
    exec(code, ns)                                  
  File "<string>", line 8, in <module>
ValueError: <jaxlib.xla_extension.ArrayImpl object at 0x562f77bb2070>

Changing chex.dataclass to dataclasses.dataclass outputs the same.

[picnixz]

What about jnp.arange? By the way, maybe it's an issue on jax then? Or matplotlib and its compatibility with jax (can you bisect the (Sphinx) commit responsibile for the issue?)

[jeertmans]

This faulty (but quite lengthy) commit is: ad61e41.

[jeertmans]

More precisely, the error can be fixed with this checkout:
git checkout 80ec9798612073be699c3fee04a1d6738303dc14 sphinx/ext/autodoc/importer.py

[jeertmans]

It's quite unexpected, to me, that the autodoc importer would set TYPE_CHECKING to True. This is what, I suppose, causes the error because jax is imported thinking the library is loaded for type checking, not actual computation purposes. Maybe that also their fault to have a different behavior when using TYPE_CHECKING too... I might create an issue on their repo!

[AA-Turner]

It's intentional, as there are often e.g. typing or collections.abc imports in the TYPE_CHECKING block, which are needed for documentation. However, it should be that failure to import with TYPE_CHECKING=True falls back to TYPE_CHECKING=False.