Untitled
taotieren
python
2 months ago
95 kB
10
No Index
❯ cat 2026-03-02T10:00:00/python-chex.log
[D 03-02 10:00:30.638 worker:220] [worker] got input: {'depend_packages': ['../python-dm-tree/python-dm-tree-0.1.9-3-x86_64.pkg.tar.zst', '../python-jax/python-jax-0.9.0-1-any.pkg.tar.zst', '../python-jaxlib/python-jaxlib-0.9.0-4-x86_64.pkg.tar.zst'], 'update_info': [['0.1.91-2', '0.1.91-2'], ['3.13', '3.14']], 'on_build_vers': [], 'commit_msg_template': 'python-chex: auto updated to {built_version}\n\nIt has been built because:\n* nvchecker detects the following updates: alpm(1): 3.13 -> 3.14\n* python-optax depends on it\n* python-underthesea depends on it\n* python-misaki depends on it', 'bindmounts': {'~/.cache/archbuild-bind-cache': '/build/.cache', '~/.cache/archbuild-bind-cache/ghcup': '/build/.ghcup', '~/.cache/pip': '/build/.cache/pip', '~/.cargo': '/build/.cargo'}, 'tmpfs': ['/build/.cache/bazel'], 'worker_no': 1, 'workerman': 'local', 'deadline': 1772420430.342692, 'reponame': 'aur-repo', 'user_agent': 'lilac/0.3 (URL; email)', 'result': '/tmp/python-chex-tuq9gbyi.lilac'}
[D 03-02 10:00:30.639 cmd:88] running ['sh', '-c', 'rm -f -- *.pkg.tar.xz *.pkg.tar.xz.sig *.pkg.tar.zst *.pkg.tar.zst.sig'], not using pty, showing output
[D 03-02 10:00:30.643 cmd:88] running ['bwrap', '--unshare-all', '--ro-bind', '/', '/', '--tmpfs', '/home', '--tmpfs', '/run', '--die-with-parent', '--tmpfs', '/tmp', '--proc', '/proc', '--dev', '/dev', '--share-net', '--bind', '/home/lilac/aur-repo/aur-repo/python-chex', '/tmp/python-chex', '--chdir', '/tmp/python-chex', '--ro-bind', PosixPath('/home/lilac/.lilac/gnupg'), '/home/lilac/.gnupg', '/bin/bash', '-c', 'source PKGBUILD && declare -p pkgver pkgrel || :'], not using pty, not showing output
[D 03-02 10:00:30.652 worker:105] oldver='0.1.91-2', newver='0.1.91-2'
[I 03-02 10:00:31.355 _client:1025] HTTP Request: GET https://aur.archlinux.org/pkgbase/python-chex "HTTP/2 303 See Other"
[I 03-02 10:00:31.931 _client:1025] HTTP Request: GET https://aur.archlinux.org/packages/python-chex "HTTP/2 200 OK"
[D 03-02 10:00:31.935 cmd:88] running ['bwrap', '--unshare-all', '--ro-bind', '/', '/', '--tmpfs', '/home', '--tmpfs', '/run', '--die-with-parent', '--tmpfs', '/tmp', '--proc', '/proc', '--dev', '/dev', '--share-net', '--bind', '/home/lilac/aur-repo/aur-repo/python-chex', '/tmp/python-chex', '--chdir', '/tmp/python-chex', '--ro-bind', PosixPath('/home/lilac/.lilac/gnupg'), '/home/lilac/.gnupg', '/bin/bash', '-c', 'source PKGBUILD && declare -p pkgver pkgrel || :'], not using pty, not showing output
[D 03-02 10:00:31.940 cmd:88] running ['git', 'ls-files'], not using pty, showing output
LICENSE
PKGBUILD
chex.diff
lilac.yaml
[I 03-02 10:00:31.941 api:471] clean directory
[D 03-02 10:00:31.941 api:476] unlink file LICENSE
[D 03-02 10:00:31.943 api:476] unlink file PKGBUILD
[D 03-02 10:00:31.943 api:476] unlink file chex.diff
[D 03-02 10:00:31.943 cmd:88] running ['git', 'reset', '--hard', 'origin/master'], not using pty, showing output
HEAD 现在位于 b9a0280 upgpg: 0.1.91-2
[D 03-02 10:00:31.948 cmd:88] running ['git', 'pull', '--no-edit'], not using pty, showing output
已经是最新的。
[D 03-02 10:00:35.561 cmd:88] running ['bwrap', '--unshare-all', '--ro-bind', '/', '/', '--tmpfs', '/home', '--tmpfs', '/run', '--die-with-parent', '--tmpfs', '/tmp', '--proc', '/proc', '--dev', '/dev', '--share-net', '--bind', '/home/lilac/aur-repo/aur-repo/python-chex', '/tmp/python-chex', '--chdir', '/tmp/python-chex', '--ro-bind', PosixPath('/home/lilac/.lilac/gnupg'), '/home/lilac/.gnupg', '/bin/bash', '-c', 'source PKGBUILD && declare -p pkgver pkgrel || :'], not using pty, not showing output
[I 03-02 10:00:35.565 api:262] pkgrel updated to 3
[D 03-02 10:00:35.566 cmd:88] running ['recv_gpg_keys'], not using pty, showing output
[D 03-02 10:00:35.574 cmd:88] running ['bwrap', '--unshare-all', '--ro-bind', '/', '/', '--tmpfs', '/home', '--tmpfs', '/run', '--die-with-parent', '--tmpfs', '/tmp', '--proc', '/proc', '--dev', '/dev', '--share-net', '--bind', '/home/lilac/aur-repo/aur-repo/python-chex', '/tmp/python-chex', '--chdir', '/tmp/python-chex', '--ro-bind', PosixPath('/home/lilac/.lilac/gnupg'), '/home/lilac/.gnupg', 'makepkg', '-od', '--noprepare', '-A'], using pty, showing output
[D 03-02 10:00:35.574 cmd:95] pty master fd=4, slave fd=5.
==> 正在创建软件包：python-chex 0.1.91-3 (2026 年03月02日 星期一 10时00分36秒)
==> 警告： 跳过依赖关系检查。
==> 获取源代码...
-> 找到 chex-0.1.91.tar.gz
==> 正在验证 source 文件，使用sha256sums...
chex-0.1.91.tar.gz ... 通过
==> 正在释放源码...
-> 正在解压缩 chex-0.1.91.tar.gz，使用 bsdtar
==> 源代码已就绪。
[D 03-02 10:00:36.128 cmd:88] running ['bwrap', '--unshare-all', '--ro-bind', '/', '/', '--tmpfs', '/home', '--tmpfs', '/run', '--die-with-parent', '--tmpfs', '/tmp', '--proc', '/proc', '--dev', '/dev', '--share-net', '--bind', '/home/lilac/aur-repo/aur-repo/python-chex', '/tmp/python-chex', '--chdir', '/tmp/python-chex', '--ro-bind', PosixPath('/home/lilac/.lilac/gnupg'), '/home/lilac/.gnupg', '/bin/bash', '-c', 'source PKGBUILD && declare -p pkgver pkgrel || :'], not using pty, not showing output
[I 03-02 10:00:36.134 api:262] pkgrel updated to 4
[D 03-02 10:00:36.648 pkgbuild:143] new package: python-chex 0.1.91-4
[I 03-02 10:00:36.648 worker:194] Running build command: ['env', 'GNUPGHOME=/home/lilac/.lilac/gnupg', 'aur-repo-x86_64-build', '--', '-I', '../python-dm-tree/python-dm-tree-0.1.9-3-x86_64.pkg.tar.zst', '-I', '../python-jax/python-jax-0.9.0-1-any.pkg.tar.zst', '-I', '../python-jaxlib/python-jaxlib-0.9.0-4-x86_64.pkg.tar.zst', '-d', '/home/lilac/.cargo:/build/.cargo', '-d', '/home/lilac/.cache/pip:/build/.cache/pip', '-d', '/home/lilac/.cache/archbuild-bind-cache/ghcup:/build/.ghcup', '-d', '/home/lilac/.cache/archbuild-bind-cache:/build/.cache', '-t', '/build/.cache/bazel', '-l', 'lilac-1', '--', '--noprogressbar', '--holdver']
==> Locking clean chroot...done
:: Synchronizing package databases...
aur-repo downloading...
core downloading...
extra downloading...
archlinuxcn downloading...
arch4edu downloading...
:: Starting full system upgrade...
there is nothing to do
==> Building in chroot for [aur-repo] (x86_64)...
==> Synchronizing chroot copy [/var/lib/archbuild/aur-repo-x86_64/root] -> [lilac-1]...done
loading packages...
resolving dependencies...
looking for conflicting packages...
Package (25) New Version Net Change
extra/abseil-cpp 20250814.1-1 6.55 MiB
extra/blas 3.12.1-2 0.74 MiB
extra/cblas 3.12.1-2 0.34 MiB
extra/gtest 1.17.0-2 1.63 MiB
extra/lapack 3.12.1-2 15.06 MiB
core/mpdecimal 4.0.1-1 0.33 MiB
core/python 3.14.3-1 72.18 MiB
extra/python-absl 2.4.0-1 1.56 MiB
extra/python-attrs 25.4.0-3 0.63 MiB
extra/python-charset-normalizer 3.4.4-2 0.51 MiB
extra/python-idna 3.11-2 1.18 MiB
extra/python-ml-dtypes 0.5.4-2 1.24 MiB
extra/python-numpy 2.4.2-1 47.60 MiB
extra/python-opt_einsum 3.4.0-2 0.92 MiB
extra/python-packaging 26.0-1 0.89 MiB
extra/python-platformdirs 4.9.3-1 0.40 MiB
extra/python-pooch 1.9.0-1 0.75 MiB
extra/python-requests 2.32.5-4 0.60 MiB
extra/python-scipy 1.17.1-1 115.60 MiB
extra/python-six 1.17.0-2 0.12 MiB
extra/python-urllib3 2.6.3-1 1.44 MiB
extra/python-wrapt 1.17.3-2 0.26 MiB
python-dm-tree 0.1.9-3 0.55 MiB
python-jax 0.9.0-1 23.43 MiB
python-jaxlib 0.9.0-4 230.05 MiB
Total Installed Size: 524.55 MiB
:: Proceed with installation? [Y/n]
checking keyring...
checking package integrity...
loading package files...
checking for file conflicts...
:: Processing package changes...
installing mpdecimal...
installing python...
Optional dependencies for python
python-setuptools: for building Python packages using tooling that is usually bundled with Python
python-pip: for installing Python packages using tooling that is usually bundled with Python
python-pipx: for installing Python software not packaged on Arch Linux
sqlite: for a default database integration [installed]
xz: for lzma [installed]
tk: for tkinter
installing python-six...
installing gtest...
Optional dependencies for gtest
python: gmock generator [installed]
installing abseil-cpp...
installing python-absl...
installing python-attrs...
installing blas...
installing cblas...
installing lapack...
installing python-numpy...
Optional dependencies for python-numpy
blas-openblas: faster linear algebra
installing python-wrapt...
installing python-dm-tree...
installing python-ml-dtypes...
installing python-platformdirs...
installing python-packaging...
installing python-charset-normalizer...
installing python-idna...
installing python-urllib3...
Optional dependencies for python-urllib3
python-brotli: Brotli support
python-brotlicffi: Brotli support
python-h2: HTTP/2 support
python-pysocks: SOCKS support
installing python-requests...
Optional dependencies for python-requests
python-chardet: alternative character encoding library
python-pysocks: SOCKS proxy support
installing python-pooch...
Optional dependencies for python-pooch
python-paramiko: for SFTP downloads
python-tqdm: for printing a download progress bar
installing python-scipy...
Optional dependencies for python-scipy
python-pillow: for image saving module
installing python-jaxlib...
installing python-opt_einsum...
installing python-jax...
Optional dependencies for python-jax
python-protobuf: Australis
python-requests: Distributed computations [installed]
python-kubernetes: Distributed support in k8s clusters
:: Running post-transaction hooks...
(1/1) Arming ConditionNeedsUpdate...
==> Making package: python-chex 0.1.91-4 (2026年03月02日 星期一 10时00分47秒)
==> Retrieving sources...
-> Found chex-0.1.91.tar.gz
==> Validating source files with sha256sums...
chex-0.1.91.tar.gz ... Passed
==> Making package: python-chex 0.1.91-4 (Mon Mar 2 02:00:50 2026)
==> Checking runtime dependencies...
==> Installing missing dependencies...
resolving dependencies...
looking for conflicting packages...
Package (2) New Version Net Change
extra/python-toolz 1.1.0-2 0.74 MiB
extra/python-typing_extensions 4.15.0-3 0.52 MiB
Total Installed Size: 1.26 MiB
:: Proceed with installation? [Y/n]
checking keyring...
checking package integrity...
loading package files...
checking for file conflicts...
:: Processing package changes...
installing python-toolz...
installing python-typing_extensions...
:: Running post-transaction hooks...
(1/1) Arming ConditionNeedsUpdate...
==> Checking buildtime dependencies...
==> Installing missing dependencies...
resolving dependencies...
looking for conflicting packages...
Package (10) New Version Net Change
extra/python-iniconfig 2.1.0-3 0.05 MiB
extra/python-pluggy 1.6.0-3 0.23 MiB
extra/python-pygments 2.19.2-3 15.30 MiB
extra/python-pyproject-hooks 1.2.0-6 0.11 MiB
extra/python-build 1.4.0-1 0.24 MiB
extra/python-cloudpickle 3.1.2-2 0.18 MiB
extra/python-flit-core 3.12.0-4 0.38 MiB
extra/python-installer 0.7.0-14 0.20 MiB
extra/python-pytest 1:8.4.2-3 4.69 MiB
extra/python-wheel 0.46.3-1 0.31 MiB
Total Installed Size: 21.67 MiB
:: Proceed with installation? [Y/n]
checking keyring...
checking package integrity...
loading package files...
checking for file conflicts...
:: Processing package changes...
installing python-pyproject-hooks...
installing python-build...
Optional dependencies for python-build
python-pip: to use as the Python package installer (default)
python-uv: to use as the Python package installer
python-virtualenv: to use virtualenv for build isolation
installing python-flit-core...
installing python-installer...
installing python-wheel...
Optional dependencies for python-wheel
python-keyring: for wheel.signatures
python-xdg: for wheel.signatures
python-setuptools: for legacy bdist_wheel subcommand
installing python-cloudpickle...
installing python-iniconfig...
installing python-pluggy...
installing python-pygments...
installing python-pytest...
:: Running post-transaction hooks...
(1/1) Arming ConditionNeedsUpdate...
==> Retrieving sources...
-> Found chex-0.1.91.tar.gz
==> WARNING: Skipping all source file integrity checks.
==> Extracting sources...
-> Extracting chex-0.1.91.tar.gz with bsdtar
==> Starting build()...
* Getting build dependencies for wheel...
* Building wheel...
Successfully built chex-0.1.91-py3-none-any.whl
==> Starting check()...
============================= test session starts ==============================
platform linux -- Python 3.14.3, pytest-8.4.2, pluggy-1.6.0 -- /usr/bin/python
cachedir: .pytest_cache
rootdir: /build/python-chex/src/chex-0.1.91
configfile: pyproject.toml
collecting ... collected 670 items
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_async_mode PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_checkify_errors PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_docstring_example PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_nested_chexification__with_jit PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_nested_chexification__without_jit PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_partial_python_fn PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_static_assertion__with_jit PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_static_assertion__without_jit PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_uninspected_checks PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTest::test_wrong_order_of_wrapping PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_jitted_main PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_jitted_nested_call_main PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_jitted_nested_call_thread PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_jitted_nested_wrap_main PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_jitted_nested_wrap_thread PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_jitted_thread PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_jitted_vmapped_main PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_jitted_vmapped_thread PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_pmapped_main PASSED
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_pmapped_thread PASSED
chex/_src/asserts_chexify_test.py::AssertsLibraryTest::test_assert_tree_all_finite PASSED
chex/_src/asserts_chexify_test.py::AssertsLibraryTest::test_assert_trees_all_close PASSED
chex/_src/asserts_chexify_test.py::AssertsLibraryTest::test_assert_trees_all_equal PASSED
chex/_src/asserts_chexify_test.py::AssertsLibraryTest::test_assert_trees_all_equal_with_prng_keys PASSED
chex/_src/asserts_chexify_test.py::AssertsLibraryTest::test_custom_message PASSED
chex/_src/asserts_internal_test.py::IsTraceableTest::test_is_traceable__with_jit PASSED
chex/_src/asserts_internal_test.py::IsTraceableTest::test_is_traceable__with_pmap SKIPPED
chex/_src/asserts_internal_test.py::ExceptionMessageFormatTest::test_format0 PASSED
chex/_src/asserts_internal_test.py::ExceptionMessageFormatTest::test_format1 PASSED
chex/_src/asserts_internal_test.py::ExceptionMessageFormatTest::test_format2 PASSED
chex/_src/asserts_internal_test.py::ExceptionMessageFormatTest::test_format3 PASSED
chex/_src/asserts_internal_test.py::ExceptionMessageFormatTest::test_format4 PASSED
chex/_src/asserts_internal_test.py::ExceptionMessageFormatTest::test_format5 PASSED
chex/_src/asserts_internal_test.py::ExceptionMessageFormatTest::test_format6 PASSED
chex/_src/asserts_internal_test.py::ExceptionMessageFormatTest::test_format7 PASSED
chex/_src/asserts_internal_test.py::JitCompatibleTest::test_api PASSED
chex/_src/asserts_test.py::AssertsSwitchTest::test_enable_disable_asserts PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_no_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_no_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_no_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_no_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_no_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_no_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_no_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_no_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_with_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_with_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_with_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_with_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_with_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_with_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_with_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_args_with_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_no_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_no_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_no_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_no_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_no_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_no_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_no_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_no_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_with_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_with_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_with_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_with_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_with_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_with_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_with_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type1_kwargs_with_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_no_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_no_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_no_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_no_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_no_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_no_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_no_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_no_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_with_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_with_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_with_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_with_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_with_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_with_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_with_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_args_with_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_no_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_no_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_no_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_no_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_no_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_no_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_no_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_no_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_with_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_with_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_with_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_with_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_with_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_with_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_with_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type2_kwargs_with_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_no_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_no_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_no_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_no_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_no_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_no_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_no_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_no_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_with_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_with_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_with_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_with_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_with_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_with_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_with_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_args_with_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_no_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_no_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_no_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_no_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_no_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_no_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_no_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_no_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_with_static_arg_max_traces_0__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_with_static_arg_max_traces_0__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_with_static_arg_max_traces_10__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_with_static_arg_max_traces_10__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_with_static_arg_max_traces_1__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_with_static_arg_max_traces_1__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_with_static_arg_max_traces_2__with_jit PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_assert_type3_kwargs_with_static_arg_max_traces_2__with_pmap SKIPPED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_incorrect_ordering PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_nested_functions PASSED
chex/_src/asserts_test.py::AssertMaxTracesTest::test_redefined_traced_function PASSED
chex/_src/asserts_test.py::ScalarAssertTest::test_scalar PASSED
chex/_src/asserts_test.py::ScalarAssertTest::test_scalar_in PASSED
chex/_src/asserts_test.py::ScalarAssertTest::test_scalar_in_excluded PASSED
chex/_src/asserts_test.py::ScalarAssertTest::test_scalar_negative PASSED
chex/_src/asserts_test.py::ScalarAssertTest::test_scalar_non_negative PASSED
chex/_src/asserts_test.py::ScalarAssertTest::test_scalar_positive PASSED
chex/_src/asserts_test.py::EqualSizeAssertTest::test_equal_size_should_fail_matrix PASSED
chex/_src/asserts_test.py::EqualSizeAssertTest::test_equal_size_should_fail_scalar_vector_matrix PASSED
chex/_src/asserts_test.py::EqualSizeAssertTest::test_equal_size_should_fail_vector_matrix PASSED
chex/_src/asserts_test.py::EqualSizeAssertTest::test_equal_size_should_pass_matrix PASSED
chex/_src/asserts_test.py::EqualSizeAssertTest::test_equal_size_should_pass_scalar_vector_matrix PASSED
chex/_src/asserts_test.py::EqualSizeAssertTest::test_equal_size_should_pass_vector_matrix PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_ellipsis_should_pass_multiple_ellipsis PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_ellipsis_should_pass_single_ellipsis PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_pytypes_pass PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_fail_some_wrong_size PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_fail_some_wrong_size_set PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_fail_wrong_common_shape PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_fail_wrong_common_shape_2 PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_fail_wrong_length_not_enough_sizes PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_fail_wrong_length_too_many_sizes PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_fail_wrong_size PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_pass_common_size_set PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_pass_matrices PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_pass_scalars PASSED
chex/_src/asserts_test.py::SizeAssertTest::test_size_should_pass_vectors PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_at_dims_should_fail_first_dim PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_at_dims_should_fail_last_dim PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_at_dims_should_fail_last_dim_neg PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_at_dims_should_fail_multiple_dims PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_at_dims_should_pass_first_and_last PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_at_dims_should_pass_first_and_last_neg PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_at_dims_should_pass_first_dim PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_at_dims_should_pass_first_few_dims PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_at_dims_should_pass_last_dim PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_prefix_should_fail_scalars PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_prefix_should_fail_vectors PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_prefix_should_pass_scalars PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_prefix_should_pass_vectors PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_should_fail_not_scalar PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_should_fail_wrong_length PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_should_fail_wrong_rank PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_should_pass_matrices PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_should_pass_scalars PASSED
chex/_src/asserts_test.py::EqualShapeAssertTest::test_equal_shape_should_pass_vectors PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_fail_inner_bad_prefix PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_fail_inner_bad_suffix PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_fail_inner_both_bad PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_fail_prefix PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_fail_short_rank_inner PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_fail_short_rank_prefix PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_fail_short_rank_suffix PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_fail_suffix PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_pass_inner_0 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_pass_inner_1 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_pass_inner_2 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_pass_prefix_0 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_pass_prefix_1 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_pass_prefix_2 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_pass_suffix_0 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_pass_suffix_1 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_ellipsis_should_pass_suffix_2 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_multiple_ellipses_prefix_and_suffix PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_pytypes_pass_shape_set PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_pytypes_pass_suffix PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_pytypes_pass_variable_shape PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_fail_some_wrong_shape PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_fail_some_wrong_shape_set PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_fail_wrong_common_shape PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_fail_wrong_common_shape_2 PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_fail_wrong_length_not_enough_shapes PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_fail_wrong_length_too_many_shapes PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_fail_wrong_rank PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_fail_wrong_shape PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_pass_common_shape_set PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_pass_matrices PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_pass_matrices_variable_shape PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_pass_scalars PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_pass_variable_common_shape PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_pass_vectors PASSED
chex/_src/asserts_test.py::ShapeAssertTest::test_shape_should_pass_vectors_common_shape PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_broadcastable0 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_broadcastable1 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_broadcastable2 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_broadcastable3 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_broadcastable4 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_broadcastable5 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_broadcastable6 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_not_broadcastable0 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_not_broadcastable1 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_not_broadcastable2 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_not_broadcastable3 PASSED
chex/_src/asserts_test.py::BroadcastAssertTest::test_shapes_are_not_broadcastable4 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_assert_rank_should_fail_sequence_wrong_1 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_assert_rank_should_fail_sequence_wrong_2 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_assert_rank_should_fail_sequence_wrong_3 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_assert_rank_should_fail_sequence_wrong_4 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_fail_array_expectations PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_fail_single_rank_1 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_fail_single_rank_2 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_fail_single_rank_3 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_fail_wrong_expectation_structure PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_fail_wrong_length_not_enough_ranks PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_fail_wrong_length_too_many_ranks PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_pass_multiple_options_seq_options_1 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_pass_multiple_options_seq_options_2 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_pass_multiple_options_single_option PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_pass_repeated_input_rank_1 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_pass_repeated_input_rank_2 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_pass_repeated_input_rank_3 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_pass_single_input_rank_1 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_pass_single_input_rank_2 PASSED
chex/_src/asserts_test.py::RankAssertTest::test_rank_should_pass_single_input_rank_3 PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_bfloat16_array__with_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_bfloat16_array__without_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_float32_array__with_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_float32_array__without_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_int8_array__with_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_int8_array__without_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_one_float_array__with_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_one_float_array__without_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_one_int_array__with_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_array_one_int_array__without_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_mixed PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_scalar_many_floats PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_scalar_many_floats_verbose PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_scalar_one_bool_as_float PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_scalar_one_bool_as_int PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_scalar_one_float PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_scalar_one_float_as_bool PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_scalar_one_int PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_scalar_one_int_as_bool PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_wrong_length_not_enough_types PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_fail_wrong_length_too_many_types PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_array_one_bool_array__with_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_array_one_bool_array__without_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_array_one_float_array__with_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_array_one_float_array__without_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_array_one_int_array__with_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_array_one_int_array__without_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_array_one_integer_array__with_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_array_one_integer_array__without_device PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_mixed PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_scalar_many_floats PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_scalar_many_floats_verbose PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_scalar_one_bool PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_scalar_one_float PASSED
chex/_src/asserts_test.py::TypeAssertTest::test_type_should_pass_scalar_one_int PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_axis_invalid PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_fail PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_gt_axis_invalid PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_gt_fail PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_gt_pass PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_gteq_axis_invalid PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_gteq_fail PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_gteq_pass PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_lt_axis_invalid PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_lt_fail PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_lt_pass PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_lteq_axis_invalid PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_lteq_fail PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_lteq_pass PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_pass PASSED
chex/_src/asserts_test.py::AxisDimensionAssertionsTest::test_assert_axis_dimension_string_tensor PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_has_only_ndarrays PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_on_device FAILED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_on_host FAILED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_sharded FAILED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_no_nones PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_mismatched_should_fail_matches_nothing PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_mismatched_should_fail_matches_y PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_mismatched_should_fail_matches_z PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_mismatched_should_pass_prefix PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_prefix_matrices PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_prefix_scalars PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_prefix_vectors PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_fail_ints PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_fail_none PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_fail_prefix PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_fail_suffix PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_pass_ints PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_pass_none PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_pass_prefix1 PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_pass_prefix2 PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_pass_prefix3 PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_pass_suffix1 PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_pass_suffix2 PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_should_pass_suffix3 PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_suffix_long_suffix PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_suffix_matching_matrices PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_suffix_matching_scalars PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_suffix_matching_vectors PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_suffix_mismatch_bad_suffix_leaf_1 PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_shape_suffix_mismatch_bad_suffix_leaf_2 PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_bfloat16 PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_fails_different_structure PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_fails_values_differ PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_passes_same_tree PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_passes_values_close_but_not_equal PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_passes_values_equal PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_ulp_bfloat16_fails_values_gt_maxulp_apart PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_ulp_bfloat16_passes_values_maxulp_apart PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_ulp_bfloat16_passes_values_within_maxulp PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_ulp_fails_values_gt_maxulp_apart PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_ulp_jittable_raises_valueerror PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_ulp_passes_same_tree PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_ulp_passes_values_equal PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_ulp_passes_values_maxulp_apart PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_close_ulp_passes_values_within_maxulp PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_dtypes PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_fail_values_close_but_not_equal PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_none PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_passes_same_tree PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_passes_values_equal PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_prng_keys PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_shapes PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_shapes_and_dtypes PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_sizes PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_strict_mode PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_structs PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_trees_all_equal_wrong_usage PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_leaf_shape_should_fail_wrong_length PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_tree_all_finite_passes_finite PASSED
chex/_src/asserts_test.py::TreeAssertionsTest::test_tree_all_finite_should_fail_inf PASSED
chex/_src/asserts_test.py::DevicesAssertTest::test_cpu_assert PASSED
chex/_src/asserts_test.py::DevicesAssertTest::test_gpu_assert PASSED
chex/_src/asserts_test.py::DevicesAssertTest::test_not_less_than0 PASSED
chex/_src/asserts_test.py::DevicesAssertTest::test_not_less_than1 PASSED
chex/_src/asserts_test.py::DevicesAssertTest::test_not_less_than2 PASSED
chex/_src/asserts_test.py::DevicesAssertTest::test_tpu_assert PASSED
chex/_src/asserts_test.py::DevicesAssertTest::test_unsupported_device PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_easy0 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_easy1 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_easy2 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_easy_with_stop_gradient0 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_easy_with_stop_gradient1 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_easy_with_stop_gradient2 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_hard0 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_hard1 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_hard2 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_hard_with_stop_gradient0 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_hard_with_stop_gradient1 PASSED
chex/_src/asserts_test.py::NumericalGradsAssertTest::test_hard_with_stop_gradient2 PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_fail_arrays PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_fail_dicts1 PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_fail_dicts2 PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_fail_dicts3 PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_fail_dicts4 PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_fail_dtypes PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_fail_lists PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_fail_lists2 PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_pass_dicts PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_pass_dtypes PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_pass_lists PASSED
chex/_src/asserts_test.py::EqualAssertionsTest::test_assert_equal_pass_on_arrays PASSED
chex/_src/asserts_test.py::IsDivisibleTest::test_assert_is_divisible PASSED
chex/_src/asserts_test.py::IsDivisibleTest::test_assert_is_divisible_fail PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testFlattenAndUnflatten_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testFlattenAndUnflatten_original PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testFlattenUpTo_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testFlattenUpTo_original PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testFlattenWithPathUpTo_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testFlattenWithPathUpTo_original PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testFlattenWithPath_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testFlattenWithPath_original PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testIsDataclass_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testIsDataclass_original PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testMapStructureUpTo_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testMapStructureUpTo_original PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testMapStructureWithPathUpTo_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testMapStructureWithPathUpTo_original PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testMapStructureWithPath_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testMapStructureWithPath_original PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testMapStructure_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testMapStructure_original PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testTraverse_chex PASSED
chex/_src/dataclass_test.py::MappableDataclassTest::testTraverse_original PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_instance_fields0 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_instance_fields1 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_mappable_dataclass_false PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_replace0 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_replace1 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_requires_kwargs_by_default PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_tree_leaves0 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_tree_leaves1 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_tree_map0 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_dataclass_tree_map1 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_decorator_alias PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_disallowed_fields PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_flatten_is_leaf0 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_flatten_is_leaf1 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_flatten_respects_post_init PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_flatten_roundtrip_ordering0 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_flatten_roundtrip_ordering1 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_frozen_dataclass_raise_error PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_generic_dataclass_mappable PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_generic_dataclass_not_mappable PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_get_and_set_state_frozen PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_get_and_set_state_mutable PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_inheritance_from_empty_frozen_base PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_inheritance_frozen PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_inheritance_is_possible_thanks_to_kw_only PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_inheritance_mutable PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_keys_and_values_override0 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_keys_and_values_override1 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_keys_and_values_type0 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_keys_and_values_type1 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_mappable_eq_override PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_roundtrip_serialization0 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_roundtrip_serialization1 PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_tree_flatten_with_keys PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_tree_map_with_keys PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_tree_map_with_keys_traversal_order PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_tuple_conversion PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_tuple_rev_conversion_frozen PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_tuple_rev_conversion_mutable PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_unexpected_kwargs PASSED
chex/_src/dataclass_test.py::DataclassesTest::test_unfrozen_dataclass_is_mutable PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_docstring_examples **********************************************************************
File "/build/python-chex/src/chex-0.1.91/chex/_src/dimensions.py", line 108, in NoName
Failed example:
w = y * x # Cannot broadcast (3, 5, 7) with (3, 5)
Expected:
Traceback (most recent call last):
...
ValueError: Incompatible shapes for broadcasting: ((3, 5, 7), (1, 3, 5))
Got:
Traceback (most recent call last):
File "/usr/lib/python3.14/site-packages/jax/_src/lax/lax.py", line 215, in _broadcast_shapes_uncached
return _try_broadcast_shapes(*rank_promoted_shapes, name='broadcast_shapes')
File "/usr/lib/python3.14/site-packages/jax/_src/lax/lax.py", line 130, in _try_broadcast_shapes
raise TypeError(f'{name} got incompatible shapes for broadcasting: '
f'{", ".join(map(str, map(tuple, shapes)))}.')
TypeError: broadcast_shapes got incompatible shapes for broadcasting: (3, 5, 7), (1, 3, 5).
<BLANKLINE>
The above exception was the direct cause of the following exception:
<BLANKLINE>
Traceback (most recent call last):
File "/usr/lib/python3.14/site-packages/jax/_src/lax/lax.py", line 193, in broadcast_shapes
return _broadcast_shapes_cached(*shapes)
File "/usr/lib/python3.14/site-packages/jax/_src/util.py", line 301, in wrapper
return cached(trace_context(), *args, **kwargs)
File "/usr/lib/python3.14/site-packages/jax/_src/util.py", line 295, in cached
return f(*args, **kwargs)
File "/usr/lib/python3.14/site-packages/jax/_src/lax/lax.py", line 199, in _broadcast_shapes_cached
return _broadcast_shapes_uncached(*shapes)
File "/usr/lib/python3.14/site-packages/jax/_src/lax/lax.py", line 218, in _broadcast_shapes_uncached
raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}") from err
ValueError: Incompatible shapes for broadcasting: shapes=[(3, 5, 7), (3, 5)]
<BLANKLINE>
During handling of the above exception, another exception occurred:
<BLANKLINE>
Traceback (most recent call last):
File "/usr/lib/python3.14/site-packages/jax/_src/lax/lax.py", line 215, in _broadcast_shapes_uncached
return _try_broadcast_shapes(*rank_promoted_shapes, name='broadcast_shapes')
File "/usr/lib/python3.14/site-packages/jax/_src/lax/lax.py", line 130, in _try_broadcast_shapes
raise TypeError(f'{name} got incompatible shapes for broadcasting: '
f'{", ".join(map(str, map(tuple, shapes)))}.')
TypeError: broadcast_shapes got incompatible shapes for broadcasting: (3, 5, 7), (1, 3, 5).
<BLANKLINE>
The above exception was the direct cause of the following exception:
<BLANKLINE>
Traceback (most recent call last):
File "<doctest NoName[15]>", line 1, in <module>
w = y * x # Cannot broadcast (3, 5, 7) with (3, 5)
~~^~~
File "/usr/lib/python3.14/site-packages/jax/_src/numpy/array_methods.py", line 609, in deferring_binary_op
return binary_op(*args)
File "/usr/lib/python3.14/site-packages/jax/_src/numpy/ufunc_api.py", line 182, in __call__
return call(*args)
File "/usr/lib/python3.14/site-packages/jax/_src/numpy/ufuncs.py", line 1280, in multiply
x, y = promote_args("multiply", x, y)
~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.14/site-packages/jax/_src/numpy/util.py", line 222, in promote_args
return promote_shapes(fun_name, *promote_dtypes(*args))
File "/usr/lib/python3.14/site-packages/jax/_src/numpy/util.py", line 60, in promote_shapes
result_rank = len(lax.broadcast_shapes(*shapes))
~~~~~~~~~~~~~~~~~~~~^^^^^^^^^
ValueError: Incompatible shapes for broadcasting: shapes=[(3, 5, 7), (3, 5)]
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_empty_end PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_empty_mid PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_empty_solo PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_empty_start PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_neg_sized_dim PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_nested PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_unmatched_closed PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_unmatched_open PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_wildcard_asterisk PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_wildcard_named PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_fail_zero_sized_dim PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_ok_all PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_ok_end PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_ok_mid PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_ok_multiple PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_ok_start PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_ok_trivial_end PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_ok_trivial_mid PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_flatten_ok_trivial_start PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_get_exception_bad_key_string PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_get_exception_bad_key_type PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_get_literals PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_get_wildcard PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_exception_bad_key_string PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_exception_bad_key_type PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_exception_iterator PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_exception_len_mismatch PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_exception_non_integer_size PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_exception_scalar PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_ok_case_sensitive PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_ok_list PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_ok_numpy_array PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_ok_scalar PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_ok_vector PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_set_wildcard PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_size_fail_wildcard_asterisk PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_size_fail_wildcard_named PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_size_fail_wildcard_negative PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_size_fail_wildcard_zero PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_size_ok_nonscalar PASSED
chex/_src/dimensions_test.py::DimensionsTest::test_size_ok_scalar PASSED
chex/_src/fake_set_n_cpu_devices_test.py::DevicesSetterTest::test_set_n_cpu_devices SKIPPED
chex/_src/fake_test.py::PmapFakeTest::test_assert_jitted PASSED
chex/_src/fake_test.py::PmapFakeTest::test_assert_pmapped PASSED
chex/_src/fake_test.py::PmapFakeTest::test_fake_jit_faked_jit PASSED
chex/_src/fake_test.py::PmapFakeTest::test_fake_jit_plain_jit PASSED
chex/_src/fake_test.py::PmapFakeTest::test_fake_parallel_axis PASSED
chex/_src/fake_test.py::PmapFakeTest::test_fake_pmap__faked_pmap PASSED
chex/_src/fake_test.py::PmapFakeTest::test_fake_pmap__faked_pmap_but_jit PASSED
chex/_src/fake_test.py::PmapFakeTest::test_fake_pmap__plain_pmap PASSED
chex/_src/fake_test.py::PmapFakeTest::test_fake_pmap__plain_pmap_but_jit PASSED
chex/_src/fake_test.py::PmapFakeTest::test_fake_pmap_axis_name PASSED
chex/_src/fake_test.py::PmapFakeTest::test_parallel_ops_equivalence PASSED
chex/_src/fake_test.py::PmapFakeTest::test_pmap_and_jit_fake_both PASSED
chex/_src/fake_test.py::PmapFakeTest::test_pmap_and_jit_fake_jit PASSED
chex/_src/fake_test.py::PmapFakeTest::test_pmap_and_jit_fake_nothing PASSED
chex/_src/fake_test.py::PmapFakeTest::test_pmap_and_jit_fake_pmap PASSED
chex/_src/fake_test.py::PmapFakeTest::test_pmap_with_complex_static_broadcasted_object0 PASSED
chex/_src/fake_test.py::PmapFakeTest::test_pmap_with_complex_static_broadcasted_object1 PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_default_params_fake_both PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_default_params_fake_jit PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_default_params_fake_nothing PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_default_params_fake_pmap PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_kwargs_fake_both PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_kwargs_fake_jit PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_kwargs_fake_nothing PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_kwargs_fake_pmap PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_partial_fake_both PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_partial_fake_jit PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_partial_fake_nothing PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_partial_fake_pmap PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_static_broadcasted_argnums_fake_nothing PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_static_broadcasted_argnums_fake_nothing_no_static_args PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_static_broadcasted_argnums_fake_pmap PASSED
chex/_src/fake_test.py::PmapFakeTest::test_with_static_broadcasted_argnums_fake_pmap_no_static_args PASSED
chex/_src/fake_test.py::OnCallOfTransformedFunctionTest::test_on_call_of_transformed_function PASSED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_detects_explicitly_forbidden_cpu_computation SKIPPED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_detects_explicitly_forbidden_tpu_computation SKIPPED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_detects_implicitly_forbidden_cpu_computation SKIPPED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_detects_implicitly_forbidden_tpu_computation SKIPPED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_ignores_explicitly_allowed_cpu_computation SKIPPED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_ignores_explicitly_allowed_tpu_computation SKIPPED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_ignores_implicitly_allowed_cpu_computation SKIPPED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_ignores_implicitly_allowed_tpu_computation SKIPPED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_raises_if_contradictory_restrictions_specified SKIPPED
chex/_src/restrict_backends_test.py::RestrictBackendsTest::test_raises_if_no_restrictions_specified SKIPPED
chex/_src/variants_test.py::ParamsProductTest::test_product PASSED
chex/_src/variants_test.py::FailedTestsTest::test_useful_errors PASSED
chex/_src/variants_test.py::FailedTestsTest::test_useful_failures PASSED
chex/_src/variants_test.py::OneFailedVariantTest::test_useful_failure PASSED
chex/_src/variants_test.py::WrongBaseClassTest::test_wrong_base_class PASSED
chex/_src/variants_test.py::BaseClassesTest::test_inheritance_parameterized PASSED
chex/_src/variants_test.py::BaseClassesTest::test_inheritance_variants PASSED
chex/_src/variants_test.py::BaseClassesTest::test_inheritance_variants_and_parameterized PASSED
chex/_src/variants_test.py::VariantsTestCaseWithParameterizedTest::test_should_pass PASSED
chex/_src/variants_test.py::WrongWrappersOrderTest::test_incorrect_wrapping_order_all_variants PASSED
chex/_src/variants_test.py::WrongWrappersOrderTest::test_incorrect_wrapping_order_named_all_variants PASSED
chex/_src/variants_test.py::WrongWrappersOrderTest::test_incorrect_wrapping_order_named_some_variants PASSED
chex/_src/variants_test.py::WrongWrappersOrderTest::test_incorrect_wrapping_order_some_variants PASSED
chex/_src/variants_test.py::UnusedVariantTest::test_unused_variant PASSED
chex/_src/variants_test.py::NoVariantsTest::test_no_variants PASSED
chex/_src/variants_test.py::UnknownVariantArgumentsTest::test_unknown_argument PASSED
chex/_src/variants_test.py::VariantTypesTest::test_consistency PASSED
chex/_src/variants_test.py::VariantTypesTest::test_var_type_fetch PASSED
chex/_src/variants_test.py::CountVariantsTest::test_counters PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants__with_device PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants__with_jit PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants__with_pmap SKIPPED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants__without_device PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants__without_jit PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants_no_parens__with_device PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants_no_parens__with_jit PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants_no_parens__with_pmap SKIPPED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants_no_parens__without_device PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_all_variants_no_parens__without_jit PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_many_variants_case_0__with_device PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_many_variants_case_0__with_jit PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_many_variants_case_0__without_device PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_many_variants_case_0__without_jit PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_many_variants_case_1__with_device PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_many_variants_case_1__with_jit PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_many_variants_case_1__without_device PASSED
chex/_src/variants_test.py::MultipleVariantsTest::test_many_variants_case_1__without_jit PASSED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_0__with_device PASSED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_0__with_jit PASSED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_0__with_pmap SKIPPED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_0__without_device PASSED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_0__without_jit PASSED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_1__with_device PASSED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_1__with_jit PASSED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_1__with_pmap SKIPPED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_1__without_device PASSED
chex/_src/variants_test.py::VmappedFunctionTest::test_vmapped_fn_named_params_case_1__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_called_variant__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_named_scalar_parameters_case_0__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_named_scalar_parameters_case_1__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_ndarray_named_parameters_case_0__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_ndarray_named_parameters_case_1__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_ndarray_parameters_(1, 2, 1)__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_ndarray_parameters_(4, 6, 2)__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_scalar_parameters_(1, 2, 1)__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_scalar_parameters_(4, 6, 2)__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_with_kwargs__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_with_ndarray_args__without_jit PASSED
chex/_src/variants_test.py::WithoutJitTest::test_with_scalar_args__without_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_called_variant__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_different_jit_kwargs_(1, 2, 1)__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_different_jit_kwargs_(4, 6, 2)__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_named_scalar_parameters_case_0__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_named_scalar_parameters_case_1__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_ndarray_named_parameters_case_0__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_ndarray_named_parameters_case_1__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_ndarray_parameters_(1, 2, 1)__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_ndarray_parameters_(4, 6, 2)__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_scalar_parameters_(1, 2, 1)__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_scalar_parameters_(4, 6, 2)__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_with_kwargs__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_with_ndarray_args__with_jit PASSED
chex/_src/variants_test.py::WithJitTest::test_with_scalar_args__with_jit PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_called_variant__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_emplace_case_0__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_emplace_case_1__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_named_scalar_parameters_case_0__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_named_scalar_parameters_case_1__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_ndarray_named_parameters_case_0__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_ndarray_named_parameters_case_1__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_ndarray_parameters_(1, 2, 1)__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_ndarray_parameters_(4, 6, 2)__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_scalar_parameters_(1, 2, 1)__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_scalar_parameters_(4, 6, 2)__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_with_kwargs__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_with_ndarray_args__without_device PASSED
chex/_src/variants_test.py::WithoutDeviceTest::test_with_scalar_args__without_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_called_variant__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_emplace_case_0__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_emplace_case_1__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_ignore_argnums_case_0__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_ignore_argnums_case_1__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_named_scalar_parameters_case_0__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_named_scalar_parameters_case_1__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_ndarray_named_parameters_case_0__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_ndarray_named_parameters_case_1__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_ndarray_parameters_(1, 2, 1)__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_ndarray_parameters_(4, 6, 2)__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_scalar_parameters_(1, 2, 1)__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_scalar_parameters_(4, 6, 2)__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_with_kwargs__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_with_ndarray_args__with_device PASSED
chex/_src/variants_test.py::WithDeviceTest::test_with_scalar_args__with_device PASSED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_called_variant__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_named_scalar_parameters_case_0__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_named_scalar_parameters_case_1__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_ndarray_named_parameters_case_0__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_ndarray_named_parameters_case_1__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_ndarray_parameters_(1, 2, 1)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_ndarray_parameters_(4, 6, 2)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_scalar_parameters_(1, 2, 1)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_scalar_parameters_(4, 6, 2)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_with_kwargs__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_with_ndarray_args__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapSingleDeviceTest::test_with_scalar_args__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapAllAvailableDeviceTest::test_pmap_(1, 2, 1)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapAllAvailableDeviceTest::test_pmap_(4, 6, 2)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapAllAvailableDeviceTest::test_pmap_in_axes__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapAllAvailableDeviceTest::test_pmap_static_argnums_(1, 2, 1)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapAllAvailableDeviceTest::test_pmap_static_argnums_(4, 6, 2)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapAllAvailableDeviceTest::test_pmap_static_argnums_zero__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapAllAvailableDeviceTest::test_pmap_vmapped_fn_(1, 2, 1)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapAllAvailableDeviceTest::test_pmap_vmapped_fn_(4, 6, 2)__with_pmap SKIPPED
chex/_src/variants_test.py::WithPmapAllAvailableDeviceTest::test_pmap_wrong_axis_size__with_pmap SKIPPED
chex/_src/warnings_test.py::WarningsTest::test_create_deprecated_function_alias PASSED
chex/_src/warnings_test.py::WarningsTest::test_warn_deprecated_function PASSED
chex/_src/warnings_test.py::WarningsTest::test_warn_deprecated_function_no_replacement PASSED
chex/_src/warnings_test.py::WarningsTest::test_warn_only_n_pos_args_in_future PASSED
chex/chex_test.py::ChexTest::test_import PASSED
=================================== FAILURES ===================================
_______________ TreeAssertionsTest.test_assert_tree_is_on_device _______________
self = <chex._src.asserts_test.TreeAssertionsTest testMethod=test_assert_tree_is_on_device>
def test_assert_tree_is_on_device(self):
# Check CPU platform.
cpu = jax.local_devices(backend='cpu')[0]
to_cpu = lambda x: jax.device_put(x, cpu)
cpu_tree = {'a': to_cpu(np.zeros(1)), 'b': to_cpu(np.ones(3))}
asserts.assert_tree_is_on_device(cpu_tree, device=cpu)
asserts.assert_tree_is_on_device(cpu_tree, platform='cpu')
asserts.assert_tree_is_on_device(cpu_tree, platform=['cpu'])
asserts.assert_tree_is_on_device(cpu_tree, device=cpu, platform='')
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'a\' resides on \'cpu\'')):
asserts.assert_tree_is_on_device(cpu_tree, platform='tpu')
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'b\' resides on \'cpu\'')):
asserts.assert_tree_is_on_device(cpu_tree, platform=('tpu', 'gpu'))
# Check TPU platform (if available).
if _num_devices_available('tpu') > 1:
tpu_1, tpu_2 = jax.devices('tpu')[:2]
to_tpu_1 = lambda x: jax.device_put(x, tpu_1)
to_tpu_2 = lambda x: jax.device_put(x, tpu_2)
tpu_1_tree = {'a': to_tpu_1(np.zeros(1)), 'b': to_tpu_1(np.ones(3))}
tpu_2_tree = {'a': to_tpu_2(np.zeros(1)), 'b': to_tpu_2(np.ones(3))}
tpu_1_2_tree = {'a': to_tpu_1(np.zeros(1)), 'b': to_tpu_2(np.ones(3))}
# Device asserts.
asserts.assert_tree_is_on_device(tpu_1_tree, device=tpu_1)
asserts.assert_tree_is_on_device(tpu_2_tree, device=tpu_2)
with self.assertRaisesRegex(
AssertionError, _get_err_regex(r"'a' resides on.*TpuDevice\(id=0")
):
asserts.assert_tree_is_on_device(tpu_1_tree, device=tpu_2)
with self.assertRaisesRegex(
AssertionError, _get_err_regex(r"'a' resides on.*TpuDevice\(id=1")
):
asserts.assert_tree_is_on_device(tpu_2_tree, device=tpu_1)
with self.assertRaisesRegex(
AssertionError, _get_err_regex("'a' resides on .*Cpu")
):
asserts.assert_tree_is_on_device(cpu_tree, device=tpu_2)
# Platform asserts.
asserts.assert_tree_is_on_device(tpu_1_tree, platform='tpu')
asserts.assert_tree_is_on_device(tpu_2_tree, platform='tpu')
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'a\' resides on \'tpu\'')):
asserts.assert_tree_is_on_device(tpu_1_tree, platform='cpu')
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'a\' resides on \'tpu\'')):
asserts.assert_tree_is_on_device(tpu_2_tree, platform='gpu')
# Mixed cases.
asserts.assert_tree_is_on_device(tpu_1_2_tree, platform='tpu')
asserts.assert_tree_is_on_device((tpu_1_2_tree, cpu_tree),
platform=('cpu', 'tpu'))
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'1/a\' resides on \'cpu\'')):
asserts.assert_tree_is_on_device((tpu_1_2_tree, cpu_tree),
platform='tpu')
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'0/a\' resides on \'tpu\'')):
asserts.assert_tree_is_on_device((tpu_1_2_tree, cpu_tree),
platform=('cpu', 'gpu'))
# Check incorrect inputs.
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'b\' is not an ndarray')):
asserts.assert_tree_is_on_device({'a': np.zeros(1), 'b': 1})
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'b\' has unexpected type')):
asserts.assert_tree_is_on_device({'a': jnp.zeros(1), 'b': np.ones(3)})
> with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'a\' is a ShardedDeviceArra')):
E AssertionError: AssertionError not raised
chex/_src/asserts_test.py:1318: AssertionError
________________ TreeAssertionsTest.test_assert_tree_is_on_host ________________
AssertionError: [Chex] Assertion assert_tree_is_on_host failed: Tree leaf 'a' resides on {CpuDevice(id=0)} (CPU devices are disallowed).
During handling of the above exception, another exception occurred:
self = <chex._src.asserts_test.TreeAssertionsTest testMethod=test_assert_tree_is_on_host>
def test_assert_tree_is_on_host(self):
cpu = jax.local_devices(backend='cpu')[0]
# Check Numpy arrays.
for flag in (False, True):
asserts.assert_tree_is_on_host({'a': np.zeros(1), 'b': np.ones(3)},
allow_cpu_device=flag)
asserts.assert_tree_is_on_host(np.zeros(4), allow_cpu_device=flag)
asserts.assert_tree_is_on_host(
jax.device_get(jax.device_put(np.ones(3))), allow_cpu_device=flag)
asserts.assert_tree_is_on_host((), allow_cpu_device=flag)
# Check DeviceArray (for platforms other than CPU).
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'a\' resides on')):
asserts.assert_tree_is_on_host({'a': jnp.zeros(1)},
allow_cpu_device=False)
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'a\' resides on')):
asserts.assert_tree_is_on_host({'a': jax.device_put(np.zeros(1))},
allow_cpu_device=False)
# Check Jax arrays on CPU.
cpu_arr = jax.device_put(np.ones(5), cpu)
asserts.assert_tree_is_on_host({'a': cpu_arr})
asserts.assert_tree_is_on_host({'a': np.zeros(1), 'b': cpu_arr})
# Check sharded Jax arrays on CPUs.
asserts.assert_tree_is_on_host(
{'a': jax.device_put_replicated(np.zeros(1), (cpu,))},
allow_cpu_device=True,
allow_sharded_arrays=True,
)
# Disallow JAX arrays on CPU.
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'a\' resides on.*CPU')):
asserts.assert_tree_is_on_host({'a': cpu_arr},
allow_cpu_device=False)
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'b\' resides on.*CPU')):
asserts.assert_tree_is_on_host({'a': np.zeros(1), 'b': cpu_arr},
allow_cpu_device=False)
# Check incorrect inputs.
with self.assertRaisesRegex(AssertionError,
_get_err_regex('\'b\' is not an ndarray')):
asserts.assert_tree_is_on_host({'a': np.zeros(1), 'b': 1})
# ShardedArrays are disallowed.
> with self.assertRaisesRegex(
AssertionError, _get_err_regex('sharded arrays are disallowed')
):
E AssertionError: "\[Chex\]\ [\s\S]*sharded arrays are disallowed" does not match "[Chex] Assertion assert_tree_is_on_host failed: Tree leaf 'a' resides on {CpuDevice(id=0)} (CPU devices are disallowed)."
chex/_src/asserts_test.py:1219: AssertionError
________________ TreeAssertionsTest.test_assert_tree_is_sharded ________________
self = <chex._src.asserts_test.TreeAssertionsTest testMethod=test_assert_tree_is_sharded>
def test_assert_tree_is_sharded(self):
np_tree = {'a': np.zeros(1), 'b': np.ones(3)}
def _format(*devs):
return re.escape(f'{devs}')
# Check single-device case.
cpu = jax.local_devices(backend='cpu')[0]
cpu_tree = jax.device_put_replicated(np_tree, (cpu,))
> asserts.assert_tree_is_sharded(cpu_tree, devices=(cpu,))
chex/_src/asserts_test.py:1334:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
chex/_src/asserts_internal.py:279: in _chex_assert_fn
host_assertion_fn(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
custom_message = None, custom_message_format_vars = ()
include_default_message = True, exception_type = <class 'AssertionError'>
args = ({'a': Array([[0.]], dtype=float32), 'b': Array([[1., 1., 1.]], dtype=float32)},)
kwargs = {'devices': (CpuDevice(id=0),)}
assertion_exc = AssertionError("Tree leaf 'a' is not sharded (devices={CpuDevice(id=0)}).\nTree leaf 'b' is not sharded (devices={CpuDevice(id=0)}).")
value_exc = None
error_msg = "[Chex] Assertion assert_tree_is_sharded failed: Tree leaf 'a' is not sharded (devices={CpuDevice(id=0)}).\nTree leaf 'b' is not sharded (devices={CpuDevice(id=0)})."
default_msg = 'Assertion assert_tree_is_sharded failed: '
def _assert_on_host(*args,
custom_message: Optional[str] = None,
custom_message_format_vars: Sequence[Any] = (),
include_default_message: bool = True,
exception_type: Type[Exception] = AssertionError,
**kwargs) -> None:
# Format error's stack trace to remove Chex' internal frames.
assertion_exc = None
value_exc = None
try:
assert_fn(*args, **kwargs)
except AssertionError as e:
assertion_exc = e
except ValueError as e:
value_exc = e
finally:
if value_exc is not None:
raise ValueError(str(value_exc))
if assertion_exc is not None:
# Format the exception message.
error_msg = str(assertion_exc)
# Include only the name of the outermost chex assertion.
if error_msg.startswith(ERR_PREFIX):
error_msg = error_msg[error_msg.find("failed:") + len("failed:"):]
# Whether to include the default error message.
default_msg = (f"Assertion {name} failed: "
if include_default_message else "")
error_msg = f"{ERR_PREFIX}{default_msg}{error_msg}"
# Whether to include a custom error message.
if custom_message:
if custom_message_format_vars:
custom_message = custom_message.format(*custom_message_format_vars)
error_msg = f"{error_msg} [{custom_message}]"
> raise exception_type(error_msg)
E AssertionError: [Chex] Assertion assert_tree_is_sharded failed: Tree leaf 'a' is not sharded (devices={CpuDevice(id=0)}).
E Tree leaf 'b' is not sharded (devices={CpuDevice(id=0)}).
chex/_src/asserts_internal.py:197: AssertionError
=============================== warnings summary ===============================
chex/_src/asserts_chexify_test.py::AssertsChexifyTestSuite::test_log_abs_fn_jitted_main
/build/python-chex/src/chex-0.1.91/chex/_src/asserts_chexify_test.py:404: DeprecationWarning: jax.device_put_replicated is deprecated; use jax.device_put instead.
(x_pos, x_with_neg) = jax.device_put_replicated((x_pos, x_with_neg),
chex/_src/asserts_test.py: 20 warnings
/build/python-chex/src/chex-0.1.91/chex/_src/asserts.py:831: DeprecationWarning: Implicit conversion of an array to a dtype is deprecated; rather than dtype=arr use dtype=arr.dtype. In the future this will result in an error.
dtype = np.result_type(x)
chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_on_device
/build/python-chex/src/chex-0.1.91/chex/_src/asserts.py:1061: DeprecationWarning: jax.sharding.PmapSharding is deprecated; use jax.sharding.NamedSharding instead.
if isinstance(x.sharding, jax.sharding.PmapSharding):
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=========================== short test summary info ============================
FAILED chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_on_device
FAILED chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_on_host
FAILED chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_sharded
====== 3 failed, 582 passed, 85 skipped, 22 warnings in 79.20s (0:01:19) =======
==> ERROR: A failure occurred in check().
Aborting...Editor is loading...
Leave a Comment