gurudesh/pytracebugs huggingface dataset:
# Buggy code
def find_spec(self, fullname, path, target=None):
    if fullname in _DISTUTILS_PATCH and self.fullname is None:
        with self.lock:
            self.fullname = fullname
            try:
                spec = find_spec(fullname, path)
                if spec is not None:
                    # https://www.python.org/dev/peps/pep-0451/#how-loading-will-work
                    is_new_api = hasattr(spec.loader, "exec_module")
                    func_name = "exec_module" if is_new_api else "load_module"
                    old = getattr(spec.loader, func_name)
                    func = self.exec_module if is_new_api else self.load_module
                    if old is not func:
                        setattr(spec.loader, func_name, partial(func, old))
                    return spec
            finally:
                self.fullname = None


# Traceback (AttributeError)
import setuptools
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "<frozen importlib._bootstrap>", line 971, in _find_and_load
File "<frozen importlib._bootstrap>", line 955, in _find_and_load_unlocked
File "<frozen importlib._bootstrap>", line 656, in _load_unlocked
File "<frozen importlib._bootstrap>", line 626, in _load_backward_compatible
File "/tmp/test/setuptools-42.0.0/setuptools-42.0.0-py3.6.egg/setuptools/__init__.py", line 20, in <module>
File "<frozen importlib._bootstrap>", line 971, in _find_and_load
File "<frozen importlib._bootstrap>", line 951, in _find_and_load_unlocked
File "<frozen importlib._bootstrap>", line 894, in _find_spec
File "/tmp/salt-tests-tmpdir/tmp8v5ucrk2/psetuptools/lib/python3.6/site-packages/_virtualenv.py", line 64, in find_spec
setattr(spec.loader, func_name, partial(func, old))
AttributeError: 'zipimport.zipimporter' object attribute 'load_module' is read-only

# Fixed code
def find_spec(self, fullname, path, target=None):
    if fullname in _DISTUTILS_PATCH and self.fullname is None:
        with self.lock:
            self.fullname = fullname
            try:
                spec = find_spec(fullname, path)
                if spec is not None:
                    # https://www.python.org/dev/peps/pep-0451/#how-loading-will-work
                    is_new_api = hasattr(spec.loader, "exec_module")
                    func_name = "exec_module" if is_new_api else "load_module"
                    old = getattr(spec.loader, func_name)
                    func = self.exec_module if is_new_api else self.load_module
                    if old is not func:
                        try:
                            setattr(spec.loader, func_name, partial(func, old))
                        except AttributeError:
                            pass  # C-Extension loaders are r/o such as zipimporter with <python 3.7
                    return spec
            finally:
                self.fullname = None
# Buggy code
def create(self):
    dirs = self.ensure_directories()
    for directory in list(dirs):
        if any(
            i
            for i in dirs
            if i is not directory and directory.parts == i.parts[: len(directory.parts)]
        ):
            dirs.remove(directory)
    for directory in sorted(dirs):
        ensure_dir(directory)

    self.set_pyenv_cfg()
    self.pyenv_cfg.write()
    true_system_site = self.enable_system_site_package
    try:
        self.enable_system_site_package = False
        for src in self._sources:
            src.run(self, self.symlinks)
    finally:
        if true_system_site != self.enable_system_site_package:
            self.enable_system_site_package = true_system_site


# Traceback (ModuleNotFoundError)
~ % source venv/bin/activate
source: no such file or directory: venv/bin/activate
~ % source venv/usr/local/opt/python@3.6.8/bin/activate
(venv) ~ % pip list
Traceback (most recent call last):
File "/Users/csalch/venv/usr/local/opt/python@3.6.8/bin/pip", line 5, in <module>
from pip._internal.cli.main import main
ModuleNotFoundError: No module named 'pip._internal.cli.main'
(venv) ~ % which python
/Users/csalch/venv/usr/local/opt/python@3.6.8/bin/python
(venv) ~ % python -m pip list
Package             Version
------------------- -------
anisble             0.1
ansible             2.9.4
appdirs             1.4.3
cffi                1.13.2
cryptography        2.8
distlib             0.3.0
filelock            3.0.12
importlib-metadata  1.5.0
importlib-resources 1.0.2
Jinja2              2.10.3
MarkupSafe          1.1.1
pip                 19.3.1
pycparser           2.19
PyYAML              5.3
ruamel.yaml         0.16.5
ruamel.yaml.clib    0.2.0
setuptools          42.0.2
six                 1.13.0
virtualenv          20.0.4
wheel               0.33.6
zipp                2.2.0
(venv) ~ % cat $(which pip)
#!/Users/csalch/venv/usr/local/opt/python@3.6.8/bin/python
# -*- coding: utf-8 -*-
import re
import sys
from pip._internal.cli.main import main
if __name__ == '__main__':
sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
sys.exit(main())
(venv) ~ %

# Fixed code
def create(self):
    dirs = self.ensure_directories()
    for directory in list(dirs):
        if any(
            i
            for i in dirs
            if i is not directory and directory.parts == i.parts[: len(directory.parts)]
        ):
            dirs.remove(directory)
    for directory in sorted(dirs):
        ensure_dir(directory)

    self.set_pyenv_cfg()
    self.pyenv_cfg.write()
    true_system_site = self.enable_system_site_package
    try:
        self.enable_system_site_package = False
        for src in self._sources:
            src.run(self, self.symlinks)
    finally:
        if true_system_site != self.enable_system_site_package:
            self.enable_system_site_package = true_system_site
    super(ViaGlobalRefVirtualenvBuiltin, self).create()
# Buggy code
def create(self):
    if self.can_be_inline:
        self.create_inline()
    else:
        self.create_via_sub_process()
        # TODO: cleanup activation scripts
    for lib in self.libs:
        ensure_dir(lib)


# Traceback (ModuleNotFoundError)
~ % source venv/bin/activate
source: no such file or directory: venv/bin/activate
~ % source venv/usr/local/opt/python@3.6.8/bin/activate
(venv) ~ % pip list
Traceback (most recent call last):
File "/Users/csalch/venv/usr/local/opt/python@3.6.8/bin/pip", line 5, in <module>
from pip._internal.cli.main import main
ModuleNotFoundError: No module named 'pip._internal.cli.main'
(venv) ~ % which python
/Users/csalch/venv/usr/local/opt/python@3.6.8/bin/python
(venv) ~ % python -m pip list
Package             Version
------------------- -------
anisble             0.1
ansible             2.9.4
appdirs             1.4.3
cffi                1.13.2
cryptography        2.8
distlib             0.3.0
filelock            3.0.12
importlib-metadata  1.5.0
importlib-resources 1.0.2
Jinja2              2.10.3
MarkupSafe          1.1.1
pip                 19.3.1
pycparser           2.19
PyYAML              5.3
ruamel.yaml         0.16.5
ruamel.yaml.clib    0.2.0
setuptools          42.0.2
six                 1.13.0
virtualenv          20.0.4
wheel               0.33.6
zipp                2.2.0
(venv) ~ % cat $(which pip)
#!/Users/csalch/venv/usr/local/opt/python@3.6.8/bin/python
# -*- coding: utf-8 -*-
import re
import sys
from pip._internal.cli.main import main
if __name__ == '__main__':
sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
sys.exit(main())
(venv) ~ %

# Fixed code
def create(self):
    if self.can_be_inline:
        self.create_inline()
    else:
        self.create_via_sub_process()

    # TODO: cleanup activation scripts

    for lib in self.libs:
        ensure_dir(lib)
    super(Venv, self).create()
# Buggy code
def _pip_config(py_executable, python_path):
    cmd = [py_executable, "-m", "pip", "config", "list"]
    config = {}
    for line in call_subprocess(
        cmd,
        show_stdout=False,
        extra_env={"PYTHONPATH": python_path, "JYTHONPATH": python_path},
        remove_from_env=["PIP_VERBOSE", "PIP_QUIET"],
    ):
        key, _, value = line.partition("=")
        if value:
            config[key] = ast.literal_eval(value)
    return config


# Traceback (OSError)
$ pip --version
pip 9.0.1 from /usr/local/lib/python2.7/dist-packages (python 2.7)

$ pip list | grep virtualenv
virtualenv (16.3.0)

$ virtualenv /tmp/venv
New python executable in /tmp/venv/bin/python
Complete output from command /tmp/venv/bin/python -m pip config list:
ERROR: unknown command "config"
----------------------------------------
Traceback (most recent call last):
File "/usr/local/bin/virtualenv", line 11, in <module>
sys.exit(main())
File "/usr/local/lib/python2.7/dist-packages/virtualenv.py", line 793, in main
symlink=options.symlink,
File "/usr/local/lib/python2.7/dist-packages/virtualenv.py", line 1087, in create_environment
install_wheel(to_install, py_executable, search_dirs, download=download)
File "/usr/local/lib/python2.7/dist-packages/virtualenv.py", line 935, in install_wheel
_install_wheel_with_search_dir(download, project_names, py_executable, search_dirs)
File "/usr/local/lib/python2.7/dist-packages/virtualenv.py", line 964, in _install_wheel_with_search_dir
config = _pip_config(py_executable, python_path)
File "/usr/local/lib/python2.7/dist-packages/virtualenv.py", line 1038, in _pip_config
remove_from_env=["PIP_VERBOSE", "PIP_QUIET"],
File "/usr/local/lib/python2.7/dist-packages/virtualenv.py", line 886, in call_subprocess
raise OSError("Command {} failed with error code {}".format(cmd_desc, proc.returncode))
OSError: Command /tmp/venv/bin/python -m pip config list failed with error code 1

# Fixed code
def _pip_config(py_executable, python_path):
    cmd = [py_executable, "-m", "pip", "config", "list"]
    config = {}
    for line in call_subprocess(
        cmd,
        show_stdout=False,
        extra_env={"PYTHONPATH": python_path, "JYTHONPATH": python_path},
        remove_from_env=["PIP_VERBOSE", "PIP_QUIET"],
        raise_on_return_code=False,
    ):
        key, _, value = line.partition("=")
        if value:
            config[key] = ast.literal_eval(value)
    return config
# Buggy code
def get_jax_enable_x64():
    """Returns the value of the flag after GoogleInit.

    We must wait until flags have been parsed (in particular for top-level
    functions decorated with jax.jit), so we delay inspecting the value
    of the jax_enable_x64 flag until JIT time.
    """
    return config.x64_enabled


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def get_jax_enable_x64():
    """Returns the value of the flag after GoogleInit.

    We must wait until flags have been parsed (in particular for top-level
    functions decorated with jax.jit), so we delay inspecting the value
    of the jax_enable_x64 flag until JIT time.
    """
    # TODO(jblespiau): Delete when jaxlib 0.1.62 is the minimal version.
    if lib._xla_extension_version >= 4:
        return config.read("jax_enable_x64")
    else:
        return config.x64_enabled
# Buggy code
def x64_enabled(self):
    if self._thread_local_state.enable_x64 is None:
        self._thread_local_state.enable_x64 = bool(self.read("jax_enable_x64"))
    return self._thread_local_state.enable_x64


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def x64_enabled(self):
    if lib._xla_extension_version >= 4:
        if lib.jax_jit.get_enable_x64() is None:
            lib.jax_jit.set_enable_x64(bool(self.read("jax_enable_x64")))
        return lib.jax_jit.get_enable_x64()
    else:
        # TODO(jakevdp): Remove when minimum jaxlib is has extension version 4
        if self._thread_local_state.enable_x64 is None:
            self._thread_local_state.enable_x64 = bool(self.read("jax_enable_x64"))
        return self._thread_local_state.enable_x64
# Buggy code
def _set_x64_enabled(self, state):
    self._thread_local_state.enable_x64 = bool(state)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def _set_x64_enabled(self, state):
    if lib._xla_extension_version >= 4:
        lib.jax_jit.set_enable_x64(bool(state))
    else:
        # TODO(jakevdp): Remove when minimum jaxlib is has extension version 4
        self._thread_local_state.enable_x64 = bool(state)
# Buggy code
def cache(max_size=4096):
    def wrap(f):
        @functools.lru_cache(max_size)
        def cached(_, *args, **kwargs):
            return f(*args, **kwargs)

        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            if jax.core.debug_state.check_leaks:
                return f(*args, **kwargs)
            else:
                return cached(bool(FLAGS.jax_enable_x64), *args, **kwargs)

        wrapper.cache_clear = cached.cache_clear
        wrapper.cache_info = cached.cache_info
        return wrapper

    return wrap


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def cache(max_size=4096):
    def wrap(f):
        @functools.lru_cache(max_size)
        def cached(_, *args, **kwargs):
            return f(*args, **kwargs)

        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            if jax.core.debug_state.check_leaks:
                return f(*args, **kwargs)
            else:
                return cached(bool(config.x64_enabled), *args, **kwargs)

        wrapper.cache_clear = cached.cache_clear
        wrapper.cache_info = cached.cache_info
        return wrapper

    return wrap
# Buggy code
def wrap(f):
    @functools.lru_cache(max_size)
    def cached(_, *args, **kwargs):
        return f(*args, **kwargs)

    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        if jax.core.debug_state.check_leaks:
            return f(*args, **kwargs)
        else:
            return cached(bool(FLAGS.jax_enable_x64), *args, **kwargs)

    wrapper.cache_clear = cached.cache_clear
    wrapper.cache_info = cached.cache_info
    return wrapper


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def wrap(f):
    @functools.lru_cache(max_size)
    def cached(_, *args, **kwargs):
        return f(*args, **kwargs)

    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        if jax.core.debug_state.check_leaks:
            return f(*args, **kwargs)
        else:
            return cached(bool(config.x64_enabled), *args, **kwargs)

    wrapper.cache_clear = cached.cache_clear
    wrapper.cache_info = cached.cache_info
    return wrapper
# Buggy code
def wrapper(*args, **kwargs):
    if jax.core.debug_state.check_leaks:
        return f(*args, **kwargs)
    else:
        return cached(bool(FLAGS.jax_enable_x64), *args, **kwargs)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def wrapper(*args, **kwargs):
    if jax.core.debug_state.check_leaks:
        return f(*args, **kwargs)
    else:
        return cached(bool(config.x64_enabled), *args, **kwargs)
# Buggy code
def memoize(f):
    @functools.lru_cache(None)
    def memoized(_, *args, **kwargs):
        return f(*args, **kwargs)

    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        return memoized(bool(FLAGS.jax_enable_x64), *args, **kwargs)

    wrapper.cache_clear = memoized.cache_clear
    wrapper.cache_info = memoized.cache_info
    return wrapper


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def memoize(f):
    @functools.lru_cache(None)
    def memoized(_, *args, **kwargs):
        return f(*args, **kwargs)

    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        return memoized(bool(config.x64_enabled), *args, **kwargs)

    wrapper.cache_clear = memoized.cache_clear
    wrapper.cache_info = memoized.cache_info
    return wrapper
# Buggy code
def wrapper(*args, **kwargs):
    return memoized(bool(FLAGS.jax_enable_x64), *args, **kwargs)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def wrapper(*args, **kwargs):
    return memoized(bool(config.x64_enabled), *args, **kwargs)
# Buggy code
def get_jax_enable_x64():
    """Returns the value of the flag after GoogleInit.

    We must wait until flags have been parsed (in particular for top-level
    functions decorated with jax.jit), so we delay inspecting the value
    of the jax_enable_x64 flag until JIT time.
    """
    return FLAGS.jax_enable_x64


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def get_jax_enable_x64():
    """Returns the value of the flag after GoogleInit.

    We must wait until flags have been parsed (in particular for top-level
    functions decorated with jax.jit), so we delay inspecting the value
    of the jax_enable_x64 flag until JIT time.
    """
    return config.x64_enabled
# Buggy code
def canonicalize_dtype(dtype):
    """Convert from a dtype to a canonical dtype based on FLAGS.jax_enable_x64."""
    if isinstance(dtype, str) and dtype == "bfloat16":
        dtype = bfloat16
    try:
        dtype = np.dtype(dtype)
    except TypeError as e:
        raise TypeError(f"dtype {dtype!r} not understood") from e

    if FLAGS.jax_enable_x64:
        return dtype
    else:
        return _dtype_to_32bit_dtype.get(dtype, dtype)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def canonicalize_dtype(dtype):
    """Convert from a dtype to a canonical dtype based on config.x64_enabled."""
    if isinstance(dtype, str) and dtype == "bfloat16":
        dtype = bfloat16
    try:
        dtype = np.dtype(dtype)
    except TypeError as e:
        raise TypeError(f"dtype {dtype!r} not understood") from e

    if config.x64_enabled:
        return dtype
    else:
        return _dtype_to_32bit_dtype.get(dtype, dtype)
# Buggy code
def enable_x64():
    """Experimental context manager to temporarily enable X64 mode.

    Usage::

      >>> import jax.numpy as jnp
      >>> with enable_x64():
      ...   print(jnp.arange(10.0).dtype)
      ...
      float64

    See Also
    --------
    jax.experimental.disable_x64 :  temporarily disable X64 mode.
    """
    _x64_state = config.FLAGS.jax_enable_x64
    config.update("jax_enable_x64", True)
    try:
        yield
    finally:
        config.update("jax_enable_x64", _x64_state)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def enable_x64():
    """Experimental context manager to temporarily enable X64 mode.

    Usage::

      >>> import jax.numpy as jnp
      >>> with enable_x64():
      ...   print(jnp.arange(10.0).dtype)
      ...
      float64

    See Also
    --------
    jax.experimental.disable_x64 :  temporarily disable X64 mode.
    """
    _x64_state = config.x64_enabled
    config._set_x64_enabled(True)
    try:
        yield
    finally:
        config._set_x64_enabled(_x64_state)
# Buggy code
def disable_x64():
    """Experimental context manager to temporarily disable X64 mode.

    Usage::

      >>> import jax.numpy as jnp
      >>> with disable_x64():
      ...   print(jnp.arange(10.0).dtype)
      ...
      float32

    See Also
    --------
    jax.experimental.enable_x64 : temporarily enable X64 mode.
    """
    _x64_state = config.FLAGS.jax_enable_x64
    config.update("jax_enable_x64", False)
    try:
        yield
    finally:
        config.update("jax_enable_x64", _x64_state)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def disable_x64():
    """Experimental context manager to temporarily disable X64 mode.

    Usage::

      >>> import jax.numpy as jnp
      >>> with disable_x64():
      ...   print(jnp.arange(10.0).dtype)
      ...
      float32

    See Also
    --------
    jax.experimental.enable_x64 : temporarily enable X64 mode.
    """
    _x64_state = config.x64_enabled
    config._set_x64_enabled(False)
    try:
        yield
    finally:
        config._set_x64_enabled(_x64_state)
# Buggy code
def dtype_to_etype(dtype):
    """Convert from dtype to canonical etype (reading FLAGS.jax_enable_x64)."""
    return xla_client.dtype_to_etype(dtypes.canonicalize_dtype(dtype))


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def dtype_to_etype(dtype):
    """Convert from dtype to canonical etype (reading config.x64_enabled)."""
    return xla_client.dtype_to_etype(dtypes.canonicalize_dtype(dtype))
# Buggy code
def memoized_fun(fun: WrappedFun, *args):
    cache = fun_caches.setdefault(fun.f, {})
    if core.debug_state.check_leaks:
        key = (
            _copy_main_traces(fun.transforms),
            fun.params,
            args,
            bool(FLAGS.jax_enable_x64),
        )
    else:
        key = (fun.transforms, fun.params, args, bool(FLAGS.jax_enable_x64))
    result = cache.get(key, None)
    if result is not None:
        ans, stores = result
        fun.populate_stores(stores)
    else:
        ans = call(fun, *args)
        cache[key] = (ans, fun.stores)

    thread_local.most_recent_entry = weakref.ref(ans)
    return ans


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

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

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def memoized_fun(fun: WrappedFun, *args):
    cache = fun_caches.setdefault(fun.f, {})
    if core.debug_state.check_leaks:
        key = (_copy_main_traces(fun.transforms), fun.params, args, config.x64_enabled)
    else:
        key = (fun.transforms, fun.params, args, config.x64_enabled)
    result = cache.get(key, None)
    if result is not None:
        ans, stores = result
        fun.populate_stores(stores)
    else:
        ans = call(fun, *args)
        cache[key] = (ans, fun.stores)

    thread_local.most_recent_entry = weakref.ref(ans)
    return ans
# Buggy code
def abstractify(x):
    return ShapedArray(np.shape(x), dtypes.result_type(x))


# Traceback (TypeError)
import objax
import jax
m = objax.nn.Sequential([objax.nn.Linear(3, 4)])

v = objax.random.uniform(((32, 3)))
jax.eval_shape(m, v)  # ShapeDtypeStruct(shape=(32, 4), dtype=float32)

d = objax.util.EasyDict(shape=(32, 3), dtype=v.dtype)
print(repr(d.shape), repr(d.dtype)) # (32, 3) dtype('float32')
print(hasattr(d, 'shape'), hasattr(d, 'dtype')) # True True

jax.eval_shape(m, d)
#
Traceback (most recent call last):
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3417, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "<ipython-input-24-d48e201250f2>", line 1, in <module>
jax.eval_shape(m, d)
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/api.py", line 2283, in eval_shape
*map(abstractify, args_flat))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/util.py", line 36, in safe_map
return list(map(f, *args))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/api.py", line 2279, in abstractify
return ShapedArray(np.shape(x), dtypes.result_type(x))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/dtypes.py", line 294, in result_type
return canonicalize_dtype(dtype(args[0]))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/dtypes.py", line 288, in dtype
return np.result_type(x)
File "<__array_function__ internals>", line 5, in result_type
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/numpy/core/_internal.py", line 64, in _usefields
names, formats, offsets, titles = _makenames_list(adict, align)
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/numpy/core/_internal.py", line 40, in _makenames_list
format = dtype(obj[0], align=align)
TypeError: data type not understood

# Fixed code
def abstractify(x):
    return ShapedArray(np.shape(x), dtype(x))
# Buggy code
def _result_type_raw(*args):
    if len(args) < 2:
        return _jax_type(args[0])
    return _least_upper_bound(*{_jax_type(arg) for arg in args})


# Traceback (TypeError)
import objax
import jax
m = objax.nn.Sequential([objax.nn.Linear(3, 4)])

v = objax.random.uniform(((32, 3)))
jax.eval_shape(m, v)  # ShapeDtypeStruct(shape=(32, 4), dtype=float32)

d = objax.util.EasyDict(shape=(32, 3), dtype=v.dtype)
print(repr(d.shape), repr(d.dtype)) # (32, 3) dtype('float32')
print(hasattr(d, 'shape'), hasattr(d, 'dtype')) # True True

jax.eval_shape(m, d)
#
Traceback (most recent call last):
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3417, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "<ipython-input-24-d48e201250f2>", line 1, in <module>
jax.eval_shape(m, d)
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/api.py", line 2283, in eval_shape
*map(abstractify, args_flat))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/util.py", line 36, in safe_map
return list(map(f, *args))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/api.py", line 2279, in abstractify
return ShapedArray(np.shape(x), dtypes.result_type(x))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/dtypes.py", line 294, in result_type
return canonicalize_dtype(dtype(args[0]))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/dtypes.py", line 288, in dtype
return np.result_type(x)
File "<__array_function__ internals>", line 5, in result_type
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/numpy/core/_internal.py", line 64, in _usefields
names, formats, offsets, titles = _makenames_list(adict, align)
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/numpy/core/_internal.py", line 40, in _makenames_list
format = dtype(obj[0], align=align)
TypeError: data type not understood

# Fixed code
def _result_type_raw(*args):
    if len(args) == 1:
        return _jax_type(args[0])
    return _least_upper_bound(*{_jax_type(arg) for arg in args})