environment dependency

[yfpeng1234]

Dear author,
Could you please provide the requirements.txt or other files indicating the dependency.
I try to creat the conda environment myself and the dependency is as follows:
# Name Version Build Channel _libgcc_mutex 0.1 main _openmp_mutex 5.1 1_gnu absl-py 2.1.0 pypi_0 pypi array-record 0.5.1 pypi_0 pypi astunparse 1.6.3 pypi_0 pypi blas 1.0 mkl bzip2 1.0.8 h5eee18b_6 ca-certificates 2024.3.11 h06a4308_0 certifi 2024.2.2 py39h06a4308_0 chardet 5.2.0 pypi_0 pypi charset-normalizer 3.3.2 pypi_0 pypi chex 0.1.86 pypi_0 pypi click 8.1.7 pypi_0 pypi clu 0.0.12 pypi_0 pypi contextlib2 21.6.0 pypi_0 pypi cuda 11.6.1 0 nvidia cuda-cccl 11.6.55 hf6102b2_0 nvidia cuda-command-line-tools 11.6.2 0 nvidia cuda-compiler 11.6.2 0 nvidia cuda-cudart 11.6.55 he381448_0 nvidia cuda-cudart-dev 11.6.55 h42ad0f4_0 nvidia cuda-cuobjdump 11.6.124 h2eeebcb_0 nvidia cuda-cupti 11.6.124 h86345e5_0 nvidia cuda-cuxxfilt 11.6.124 hecbf4f6_0 nvidia cuda-driver-dev 11.6.55 0 nvidia cuda-gdb 12.4.127 0 nvidia cuda-libraries 11.6.1 0 nvidia cuda-libraries-dev 11.6.1 0 nvidia cuda-memcheck 11.8.86 0 nvidia cuda-nsight 12.4.127 0 nvidia cuda-nsight-compute 12.4.1 0 nvidia cuda-nvcc 11.6.124 hbba6d2d_0 nvidia cuda-nvdisasm 12.4.127 0 nvidia cuda-nvml-dev 11.6.55 haa9ef22_0 nvidia cuda-nvprof 12.4.127 0 nvidia cuda-nvprune 11.6.124 he22ec0a_0 nvidia cuda-nvrtc 11.6.124 h020bade_0 nvidia cuda-nvrtc-dev 11.6.124 h249d397_0 nvidia cuda-nvtx 11.6.124 h0630a44_0 nvidia cuda-nvvp 12.4.127 0 nvidia cuda-runtime 11.6.1 0 nvidia cuda-samples 11.6.101 h8efea70_0 nvidia cuda-sanitizer-api 12.4.127 0 nvidia cuda-toolkit 11.6.1 0 nvidia cuda-tools 11.6.1 0 nvidia cuda-visual-tools 11.6.1 0 nvidia dm-tree 0.1.8 pypi_0 pypi einops 0.8.0 pypi_0 pypi etils 1.5.2 pypi_0 pypi ffmpeg 4.3 hf484d3e_0 pytorch fire 0.6.0 pypi_0 pypi flatbuffers 24.3.25 pypi_0 pypi flax 0.8.3 pypi_0 pypi freetype 2.12.1 h4a9f257_0 fsspec 2024.3.1 pypi_0 pypi gast 0.5.4 pypi_0 pypi gds-tools 1.9.1.3 0 nvidia gmp 6.2.1 h295c915_3 gnutls 3.6.15 he1e5248_0 google-pasta 0.2.0 pypi_0 pypi grpcio 1.63.0 pypi_0 pypi h5py 3.11.0 pypi_0 pypi idna 3.7 py39h06a4308_0 importlib-metadata 7.1.0 pypi_0 pypi importlib-resources 6.4.0 pypi_0 pypi intel-openmp 2023.1.0 hdb19cb5_46306 jax 0.4.26 pypi_0 pypi jaxlib 0.4.26+cuda12.cudnn89 pypi_0 pypi jpeg 9e h5eee18b_1 keras 3.3.3 pypi_0 pypi lame 3.100 h7b6447c_0 lcms2 2.12 h3be6417_0 ld_impl_linux-64 2.38 h1181459_1 lerc 3.0 h295c915_0 libclang 18.1.1 pypi_0 pypi libcublas 11.9.2.110 h5e84587_0 nvidia libcublas-dev 11.9.2.110 h5c901ab_0 nvidia libcufft 10.7.1.112 hf425ae0_0 nvidia libcufft-dev 10.7.1.112 ha5ce4c0_0 nvidia libcufile 1.9.1.3 0 nvidia libcufile-dev 1.9.1.3 0 nvidia libcurand 10.3.5.147 0 nvidia libcurand-dev 10.3.5.147 0 nvidia libcusolver 11.3.4.124 h33c3c4e_0 nvidia libcusparse 11.7.2.124 h7538f96_0 nvidia libcusparse-dev 11.7.2.124 hbbe9722_0 nvidia libdeflate 1.17 h5eee18b_1 libffi 3.4.4 h6a678d5_1 libgcc-ng 11.2.0 h1234567_1 libgomp 11.2.0 h1234567_1 libiconv 1.16 h5eee18b_3 libidn2 2.3.4 h5eee18b_0 libnpp 11.6.3.124 hd2722f0_0 nvidia libnpp-dev 11.6.3.124 h3c42840_0 nvidia libnvjpeg 11.6.2.124 hd473ad6_0 nvidia libnvjpeg-dev 11.6.2.124 hb5906b9_0 nvidia libpng 1.6.39 h5eee18b_0 libstdcxx-ng 11.2.0 h1234567_1 libtasn1 4.19.0 h5eee18b_0 libtiff 4.5.1 h6a678d5_0 libunistring 0.9.10 h27cfd23_0 libwebp-base 1.3.2 h5eee18b_0 lz4-c 1.9.4 h6a678d5_1 markdown 3.6 pypi_0 pypi markdown-it-py 3.0.0 pypi_0 pypi markupsafe 2.1.5 pypi_0 pypi mdurl 0.1.2 pypi_0 pypi mkl 2023.1.0 h213fc3f_46344 mkl-service 2.4.0 py39h5eee18b_1 mkl_fft 1.3.8 py39h5eee18b_0 mkl_random 1.2.4 py39hdb19cb5_0 ml-collections 0.1.1 pypi_0 pypi ml-dtypes 0.3.2 pypi_0 pypi msgpack 1.0.8 pypi_0 pypi namex 0.0.8 pypi_0 pypi ncurses 6.4 h6a678d5_0 nest-asyncio 1.6.0 pypi_0 pypi nettle 3.7.3 hbbd107a_1 nsight-compute 2024.1.1.4 0 nvidia numpy 1.26.4 py39h5f9d8c6_0 numpy-base 1.26.4 py39hb5e798b_0 nvidia-cublas-cu12 12.3.4.1 pypi_0 pypi nvidia-cuda-cupti-cu12 12.3.101 pypi_0 pypi nvidia-cuda-nvcc-cu12 12.3.107 pypi_0 pypi nvidia-cuda-nvrtc-cu12 12.3.107 pypi_0 pypi nvidia-cuda-runtime-cu12 12.3.101 pypi_0 pypi nvidia-cudnn-cu12 8.9.7.29 pypi_0 pypi nvidia-cufft-cu12 11.0.12.1 pypi_0 pypi nvidia-curand-cu12 10.3.4.107 pypi_0 pypi nvidia-cusolver-cu12 11.5.4.101 pypi_0 pypi nvidia-cusparse-cu12 12.2.0.103 pypi_0 pypi nvidia-nccl-cu12 2.19.3 pypi_0 pypi nvidia-nvjitlink-cu12 12.3.101 pypi_0 pypi openh264 2.1.1 h4ff587b_0 openjpeg 2.4.0 h3ad879b_0 openssl 3.0.13 h7f8727e_1 opt-einsum 3.3.0 pypi_0 pypi optax 0.2.2 pypi_0 pypi optree 0.11.0 pypi_0 pypi orbax-checkpoint 0.5.10 pypi_0 pypi packaging 24.0 pypi_0 pypi pillow 10.3.0 py39h5eee18b_0 pip 23.3.1 py39h06a4308_0 promise 2.3 pypi_0 pypi protobuf 3.20.3 pypi_0 pypi psutil 5.9.8 pypi_0 pypi pygments 2.17.2 pypi_0 pypi python 3.9.19 h955ad1f_0 pytorch 1.13.1 py3.9_cuda11.6_cudnn8.3.2_0 pytorch pytorch-cuda 11.6 h867d48c_1 pytorch pytorch-mutex 1.0 cuda pytorch pyyaml 6.0.1 pypi_0 pypi readline 8.2 h5eee18b_0 requests 2.31.0 py39h06a4308_1 rich 13.7.1 pypi_0 pypi scipy 1.13.0 pypi_0 pypi setuptools 68.2.2 py39h06a4308_0 six 1.16.0 pypi_0 pypi sqlite 3.45.3 h5eee18b_0 tbb 2021.8.0 hdb19cb5_0 tensorboard 2.16.2 pypi_0 pypi tensorboard-data-server 0.7.2 pypi_0 pypi tensorflow 2.16.1 pypi_0 pypi tensorflow-datasets 4.9.3 pypi_0 pypi tensorflow-io-gcs-filesystem 0.37.0 pypi_0 pypi tensorflow-metadata 1.15.0 pypi_0 pypi tensorstore 0.1.58 pypi_0 pypi termcolor 2.4.0 pypi_0 pypi tk 8.6.12 h1ccaba5_0 toml 0.10.2 pypi_0 pypi toolz 0.12.1 pypi_0 pypi torchaudio 0.13.1 py39_cu116 pytorch torchvision 0.14.1 py39_cu116 pytorch tqdm 4.66.2 pypi_0 pypi typing-extensions 4.11.0 pypi_0 pypi typing_extensions 4.9.0 py39h06a4308_1 tzdata 2024a h04d1e81_0 urllib3 2.2.1 pypi_0 pypi werkzeug 3.0.2 pypi_0 pypi wheel 0.41.2 py39h06a4308_0 wrapt 1.16.0 pypi_0 pypi xz 5.4.6 h5eee18b_1 zipp 3.18.1 pypi_0 pypi zlib 1.2.13 h5eee18b_1 zstd 1.5.5 hc292b87_1
But when I run the code, it encountered the jax error as follows:
`jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/zlihm/yfp/RCIG/distill_dataset.py", line 282, in
fire.Fire(main)
File "/home/zlihm/anaconda3/envs/tf/lib/python3.9/site-packages/fire/core.py", line 143, in Fire
component_trace = _Fire(component, args, parsed_flag_args, context, name)
File "/home/zlihm/anaconda3/envs/tf/lib/python3.9/site-packages/fire/core.py", line 477, in _Fire
component, remaining_args = _CallAndUpdateTrace(
File "/home/zlihm/anaconda3/envs/tf/lib/python3.9/site-packages/fire/core.py", line 693, in _CallAndUpdateTrace
component = fn(*varargs, **kwargs)
File "/home/zlihm/yfp/RCIG/distill_dataset.py", line 230, in main
inner_result, _ = algorithms.run_rcig(coreset_images, coreset_labels, model_for_train.init, model_for_train.apply, ds_train, alg_config, key, inner_learning_rate, hvp_learning_rate, lr_tune = True)
File "/home/zlihm/yfp/RCIG/algorithms.py", line 87, in run_rcig
new_train_state, key = get_new_train_state(key, alg_config.pool_learning_rate, alg_config.model_depth, alg_config.has_bn, alg_config.linearize, net_forward_apply, net_forward_init, coreset_images_init.shape, naive_loss = alg_config.naive_loss)
File "/home/zlihm/yfp/RCIG/algorithms.py", line 50, in get_new_train_state
new_params = new_params.unfreeze()
AttributeError: 'dict' object has no attribute 'unfreeze'`

[yolky]

Hi, the code was using an older version of flax/jax which used frozen dicts. You can try comment out the ".unfreeze()" part.

[yfpeng1234]

okay, thanks a lot