jax_config.define_bool_state('flax_' + name, default, help) AttributeError: 'Config' object has no attribute 'define_bool_state'

[ChristopherBottomsOMRF]

I get this runtime error while running all vs all.

I assume that I'm using the wrong version of Jax or something?

create_ppi_all_vs_all.sh broken_1RWY.fa $HHBLITS 0.5 outdir

Writing fastas of each sequence to results_speed_ppi/fasta/
1RWY_seq1_60
Creating MSA for 1RWY_seq1_60
- 22:05:19.248 INFO: Search results will be written to results_speed_ppi/fasta//1RWY_seq1_60.hhr

- 22:05:22.577 INFO: Searching 15161831 column state sequences.
.......
498 tr|T0QX01|T0QX01_9STRA Unchara  96.4 0.00036 1.3E-09   56.6   0.0   43    1-46    249-291 (499)
499 tr|A0A2M4CI03|A0A2M4CI03_ANODA  96.4 0.00037 1.3E-09   52.2   0.0   44    1-47    181-224 (225)
500 tr|A0A183L6V4|A0A183L6V4_9TREM  96.4 0.00041 1.3E-09   44.2   0.0   20   27-46     36-55  (67)

Running pred 1 out of 2                                                            
/envs/speed_ppi/lib/python3.12/site-packages/Bio/Data/SCOPData.py:18: BiopythonDeprecationWarning: The 'Bio.Data.SCOPData' module will be deprecated in a future release of Biopython in favor of 'Bio.Data.PDBData. 
  warnings.warn(                                                                   
Traceback (most recent call last):                                                 
  File "/envs/speed_ppi/SpeedPPI/src/run_alphafold_all_vs_all.py", line 36, in <module>
    from alphafold.model import data                                               
  File "/envs/speed_ppi/SpeedPPI/src/alphafold/model/data.py", line 20, in <module>
    from alphafold.model import utils                                              
  File "/envs/speed_ppi/SpeedPPI/src/alphafold/model/utils.py", line 21, in <module>
    import haiku as hk                                                             
  File "/envs/speed_ppi/lib/python3.12/site-packages/haiku/__init__.py", line 20, in <module>
    from haiku import experimental                                                 
  File "/envs/speed_ppi/lib/python3.12/site-packages/haiku/experimental/__init__.py", line 55, in <module>
    from haiku.experimental import flax                                            
  File "/envs/speed_ppi/lib/python3.12/site-packages/haiku/experimental/flax.py", line 18, in <module>
    from haiku._src.flax.flax_module import Module                                 
  File "/envs/speed_ppi/lib/python3.12/site-packages/haiku/_src/flax/flax_module.py", line 20, in <module>
    import flax.core                                                               
  File "/envs/speed_ppi/lib/python3.12/site-packages/flax/__init__.py", line 19, in <module>
    from .configurations import (                                                  
  File "/envs/speed_ppi/lib/python3.12/site-packages/flax/configurations.py", line 93, in <module>
    flax_filter_frames = define_bool_state(                                        
                         ^^^^^^^^^^^^^^^^^^                                        
  File "/envs/speed_ppi/lib/python3.12/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'               

[patrickbryant1]

Hi,

Try this:

pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install --upgrade dm-haiku

[ChristopherBottomsOMRF]

Sorry, actually, the issue remains. I'd run into another issue that preceded this one (i.e. I had to update the version for tf-nightly in speed_ppi.yml to 2.16.0).

[Renske27]

I also try to run all vs all and get the same issue: AttributeError: 'Config' object has no attribute 'define_bool_state'
I already tried:
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install --upgrade dm-haiku

But unfortunately, the issue remains. Does someone know how to solve this?

[li210]

I also try to run all vs all and get the same issue: AttributeError: 'Config' object has no attribute 'define_bool_state' I already tried: pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html pip install --upgrade dm-haiku

But unfortunately, the issue remains. Does someone know how to solve this?

Downgrading Jax helped me (jax 0.4.23, jaxlib 0.4.23+cuda12.cudnn89).

[patrickbryant1]

You can also try: pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html