jax 0.8.2 incpatibility

[GaetanLepage]

When run against jax 0.8.2, the following tests are failing:

=========================== short test summary info ============================
FAILED chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_on_device - AssertionError: AssertionError not raised
FAILED chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_on_host - AssertionError: "\[Chex\]\ [\s\S]*sharded arrays are disallowed" does not m...
FAILED chex/_src/asserts_test.py::TreeAssertionsTest::test_assert_tree_is_sharded - AssertionError: [Chex] Assertion assert_tree_is_sharded failed: Tree leaf '...
============ 3 failed, 582 passed, 85 skipped, 2 warnings in 30.89s ============

• TreeAssertionsTest.test_assert_tree_is_on_device

>     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

      # 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

>         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

[Vivekk0712]

Hi! @GaetanLepage
I can reproduce this issue with JAX 0.8.2 on CPU-only environments. The problem is in the _check_sharding() function which needs updating for the new single-device array behavior.

Can this be assigned to me? I'd like to submit a fix.

cc: @danielsuo @hbq1