Outdated jax config

[landmanbester]

• Codex Africanus version: latest
• Python version: 3.10
• Operating System: ubuntu20.02

Description

As reported by @PNMwaniki here landmanbester/spimple#16 (comment) the following fails

Python 3.10.13 (main, Aug 25 2023, 13:20:03) [GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import africanus.util.jax_init
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/bester/.venv/testcodex/lib/python3.10/site-packages/africanus/__init__.py", line 10, in <module>
    import africanus.util.jax_init  # noqa
  File "/home/bester/.venv/testcodex/lib/python3.10/site-packages/africanus/util/jax_init.py", line 14, in <module>
    jax.config.update("jax_enable_x64", True)
AttributeError: module 'jax.config' has no attribute 'update'
>>>

Probably just a change in jax. Will take a look after the weekend

[landmanbester]

@sjperkins i) thanks for taking care of this and ii) please can you make a release with this fix its tripping me up in multiple places

[sjperkins]

The 0.3.6 release check is currently running, should be released on pypi when the succeeds.

https://github.com/ratt-ru/codex-africanus/actions/runs/8686683566