'Config' object error during import

[shahrozeabbas]

Hi,

Thanks for creating the tool! I am excited to use it but have a problem during the import step. I am using a Google Colab notebook with an A100 GPU. The package installs successfully using pip install vivs. However the following error occurs when running import vivs:

AttributeError                            Traceback (most recent call last)
[<ipython-input-3-68dd4ae6eda3>](https://localhost:8080/#) in <cell line: 1>()
----> 1 import vivs
      2 import scanpy

4 frames
[/usr/local/lib/python3.10/dist-packages/vivs/__init__.py](https://localhost:8080/#) in <module>
      1 from ._constants import REGISTRY_KEYS
      2 from ._utils import select_architecture, select_genes
----> 3 from ._vivs import VIVS
      4 
      5 __all__ = ["VIVS", "select_genes", "select_architecture", "REGISTRY_KEYS"]

[/usr/local/lib/python3.10/dist-packages/vivs/_vivs.py](https://localhost:8080/#) in <module>
      2 
      3 import fastcluster
----> 4 import flax.linen as nn
      5 import jax
      6 import jax.numpy as jnp

[/usr/local/lib/python3.10/dist-packages/flax/__init__.py](https://localhost:8080/#) in <module>
     17 # pyformat: disable
     18 
---> 19 from .configurations import (
     20     config as config,
     21 )

[/usr/local/lib/python3.10/dist-packages/flax/configurations.py](https://localhost:8080/#) in <module>
     91 flax_lazy_rng = static_bool_env('FLAX_LAZY_RNG', True)
     92 
---> 93 flax_filter_frames = define_bool_state(
     94     name='filter_frames',
     95     default=True,

[/usr/local/lib/python3.10/dist-packages/flax/configurations.py](https://localhost:8080/#) in define_bool_state(name, default, help)
     40   on runtime with `flax.config.update('flax_<config_name>', <value>)`.
     41   """
---> 42   return jax_config.define_bool_state('flax_' + name, default, help)
     43 
     44 

AttributeError: 'Config' object has no attribute 'define_bool_state'

Appreciate any input and look forward to trying!

Thanks.