KeyError raised after burn-in MCMC steps

[Yangcn1995]

Dear authors Currently I meet some problems when I run ferminet, when burn-in MCMC steps completed, it raises an error occurred in kfac_jax/_src/tag_graph_matcher.py, line 674. It is 'KeyError: a', the detail is in attached file 'error.json', please check it.

error.json

Then we try to print 'env' and 'var', we find the keys in variable 'env' changed, that is the reason we meet the error, the detail is in attached file 'error-print.json', please check it.

error-print.json

Could you please give us any advice to solve this problem?

Here is our environment:

channels:

• https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
• https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/
• https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
• https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
• https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
• defaults dependencies:
• _libgcc_mutex=0.1=conda_forge
• _openmp_mutex=4.5=2_gnu
• bzip2=1.0.8=hd590300_5
• ca-certificates=2023.11.17=hbcca054_0
• ld_impl_linux-64=2.40=h41732ed_0
• libexpat=2.5.0=hcb278e6_1
• libffi=3.4.2=h7f98852_5
• libgcc-ng=13.2.0=h807b86a_4
• libgomp=13.2.0=h807b86a_4
• libnsl=2.0.1=hd590300_0
• libsqlite=3.44.2=h2797004_0
• libuuid=2.38.1=h0b41bf4_0
• libxcrypt=4.4.36=hd590300_1
• libzlib=1.2.13=hd590300_5
• ncurses=6.4=h59595ed_2
• openssl=3.2.0=hd590300_1
• pip=23.3.2=pyhd8ed1ab_0
• python=3.11.7=hab00c5b_1_cpython
• readline=8.2=h8228510_1
• tk=8.6.13=noxft_h4845f30_101
• wheel=0.42.0=pyhd8ed1ab_0
• xz=5.2.6=h166bdaf_0
• pip:
  • absl-py==1.4.0
  • array-record==0.5.0
  • astunparse==1.6.3
  • attrs==23.2.0
  • blosc2==2.5.1
  • cachetools==5.3.2
  • certifi==2023.11.17
  • charset-normalizer==3.3.2
  • chex==0.1.6
  • click==8.1.7
  • cloudpickle==3.0.0
  • contextlib2==21.6.0
  • decorator==5.1.1
  • distrax==0.0.3
  • dm-haiku==0.0.11
  • dm-tree==0.1.8
  • etils==1.6.0
  • flatbuffers==23.5.26
  • flax==0.8.0
  • folx==0.2.2.post1
  • fsspec==2023.12.2
  • gast==0.5.4
  • google-auth==2.27.0
  • google-auth-oauthlib==1.2.0
  • google-pasta==0.2.0
  • googleapis-common-protos==1.62.0
  • grpcio==1.60.0
  • h5py==3.10.0
  • idna==3.6
  • immutabledict==4.1.0
  • importlib-resources==6.1.1
  • iniconfig==2.0.0
  • jax==0.4.23
  • jaxlib==0.3.25
  • jaxline==0.0.8
  • jaxtyping==0.2.25
  • jmp==0.0.4
  • keras==2.15.0
  • kfac-jax==0.0.2
  • libclang==16.0.6
  • markdown==3.5.2
  • markdown-it-py==3.0.0
  • markupsafe==2.1.4
  • mdurl==0.1.2
  • ml-collections==0.1.1
  • ml-dtypes==0.2.0
  • msgpack==1.0.7
  • ndindex==1.7
  • nest-asyncio==1.6.0
  • numexpr==2.9.0
  • numpy==1.26.3
  • oauthlib==3.2.2
  • opt-einsum==3.3.0
  • optax==0.0.5
  • orbax-checkpoint==0.4.4
  • packaging==23.2
  • pandas==2.2.0
  • pluggy==1.4.0
  • promise==2.3
  • protobuf==3.20.3
  • psutil==5.9.8
  • py-cpuinfo==9.0.0
  • pyasn1==0.5.1
  • pyasn1-modules==0.3.0
  • pyblock==0.6
  • pygments==2.17.2
  • pyscf==2.4.0
  • pytest==8.0.0
  • python-dateutil==2.8.2
  • pytz==2023.3.post1
  • pyyaml==6.0.1
  • requests==2.31.0
  • requests-oauthlib==1.3.1
  • rich==13.7.0
  • rsa==4.9
  • scipy==1.12.0
  • setuptools==69.0.3
  • six==1.16.0
  • tables==3.9.2
  • tabulate==0.9.0
  • tensorboard==2.15.1
  • tensorboard-data-server==0.7.2
  • tensorflow==2.15.0.post1
  • tensorflow-datasets==4.9.4
  • tensorflow-estimator==2.15.0
  • tensorflow-io-gcs-filesystem==0.35.0
  • tensorflow-metadata==1.14.0
  • tensorflow-probability==0.23.0
  • tensorstore==0.1.45
  • termcolor==2.4.0
  • toml==0.10.2
  • toolz==0.12.1
  • tqdm==4.66.1
  • typeguard==2.13.3
  • typing-extensions==4.9.0
  • tzdata==2023.4
  • urllib3==2.1.0
  • werkzeug==3.0.1
  • wrapt==1.14.1
  • zipp==3.17.0

Regards Qinmeng

[jsspencer]

You are using kfac-jax 0.0.2 from 2022 and jax 0.4.23 (from Dec 2023). The two are almost certainly not compatible. Please try with a more recent version of kfac-jax.