jax 0.8.2 incpatibility

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
```