AttributeError: module 'kfac_jax._src.utils' has no attribute 'Array'

[weizd]

/home/test2/anaconda3/envs/ferminet39/bin/python /home/test2/ferminet/ferminet/train.py
Traceback (most recent call last):
File "/home/test2/ferminet/ferminet/train.py", line 25, in
from ferminet import checkpoint
File "/home/test2/ferminet/ferminet/checkpoint.py", line 24, in
from ferminet import networks
File "/home/test2/ferminet/ferminet/networks.py", line 21, in
from ferminet import envelopes
File "/home/test2/ferminet/ferminet/envelopes.py", line 21, in
from ferminet import curvature_tags_and_blocks
File "/home/test2/ferminet/ferminet/curvature_tags_and_blocks.py", line 25, in
Array = kfac_jax.utils.Array
AttributeError: module 'kfac_jax._src.utils' has no attribute 'Array'

what should I do ? my python=3.9 , the version is right?

[jsspencer]

Please check your KFAC version; we typically support the latest commit in the kfac_jax repository. I've just pushed a small update which handles a recent kfac_jax API change.