bug: Installation issues in Python 3.11, Linux Ubuntu 22.04

[charlesll]

Bug Report

Fortuna version:

0.1.45

Current behavior:

• On Ubuntu 22.04, after creating a new Conda environment with Python 3.11, and installation with pip install aws-fortuna (first install): Import of fortuna fails. This is apparently related to an issue with Flax and Jax versions.

Expected behavior:

Loading aws-fortuna

Steps to reproduce:

In a terminal, type the commands:

$ conda create --name test python=3.11
$ conda activate test
$ pip install aws-fortuna
$ python
>>> import fortuna

Related code:

Python 3.11.9 (main, Apr 19 2024, 16:48:06) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import fortuna
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/charles/miniconda3/envs/test/lib/python3.11/site-packages/fortuna/__init__.py", line 1, in <module>
    import flax
  File "/home/charles/miniconda3/envs/test/lib/python3.11/site-packages/flax/__init__.py", line 18, in <module>
    from .configurations import (
  File "/home/charles/miniconda3/envs/test/lib/python3.11/site-packages/flax/configurations.py", line 92, in <module>
    flax_filter_frames = define_bool_state(
                         ^^^^^^^^^^^^^^^^^^
  File "/home/charles/miniconda3/envs/test/lib/python3.11/site-packages/flax/configurations.py", line 42, in define_bool_state
    return jax_config.define_bool_state('flax_' + name, default, help)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'Config' object has no attribute 'define_bool_state'
>>> 

Other information:

On Linux, upgrading Jax and Flax to the latest versions actually solved the issue. I don't know if it will create further errors though...

[benHeid]

I am experiencing the same error on MacOS.

[wistuba]

Thanks for pointing this out! Would you be willing to prepare a fix for this?

[mdhanna]

Can confirm that upgrading flax to 0.8.5 fixed this issue for me on MacOS.

[wistuba]

197 should fix this. Does anyone have time to confirm?