add requirements.txt

models/networks_flax.py

@@ -187,6 +187,7 @@ def __call__(self, inputs, output_layer_indx=-1):
 
         U = nn.Dense(features=self.layers[0], kernel_init=self.kernel_init(True), name=f'transformerU_{self.tag}')(inputs)
         V = nn.Dense(features=self.layers[0], kernel_init=self.kernel_init(True), name=f'transformerV_{self.tag}')(inputs)
+        # NOTE: if the models from https://doi.org/10.11583/DTU.24812004 are used, please comment out the following two lines
         U = self.activation(U)
         V = self.activation(V)
 

requirements.txt

@@ -0,0 +1,105 @@
+absl-py==1.4.0
+aiohttp==3.8.4
+aiosignal==1.3.1
+async-timeout==4.0.2
+attrs==23.1.0
+cached-property==1.5.2
+cachetools==5.3.1
+certifi==2023.5.7
+charset-normalizer==3.1.0
+chex==0.1.7
+cmake==3.26.3
+contourpy==1.0.7
+cycler==0.11.0
+distlib==0.3.6
+dm-tree==0.1.8
+etils==1.3.0
+filelock==3.8.0
+flax==0.6.10
+fonttools==4.39.4
+frozenlist==1.3.3
+fsspec==2023.5.0
+google-auth==2.19.0
+google-auth-oauthlib==1.0.0
+graphviz==0.20.1
+grpcio==1.54.2
+gviz-api==1.10.0
+h5py==3.8.0
+idna==3.4
+importlib-resources==5.12.0
+jax==0.4.10
+jaxlib==0.4.10+cuda12.cudnn88
+Jinja2==3.1.2
+joblib==1.2.0
+kiwisolver==1.4.4
+lightning-utilities==0.8.0
+lit==16.0.5
+Markdown==3.4.3
+markdown-it-py==2.2.0
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+mdurl==0.1.2
+meshio==5.3.4
+ml-dtypes==0.1.0
+mpmath==1.3.0
+msgpack==1.0.5
+multidict==6.0.4
+nest-asyncio==1.5.6
+networkx==3.1
+numpy==1.24.3
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+oauthlib==3.2.2
+opt-einsum==3.3.0
+optax==0.1.5
+orbax-checkpoint==0.2.3
+packaging==23.1
+pandas==2.0.2
+Pillow==9.5.0
+platformdirs==2.5.2
+protobuf==4.23.2
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pyDOE2==1.3.0
+pydot==1.4.2
+Pygments==2.15.1
+pyparsing==3.0.9
+python-dateutil==2.8.2
+pytorch-lightning==2.0.2
+pytz==2023.3
+PyYAML==6.0
+requests==2.31.0
+requests-oauthlib==1.3.1
+rich==13.4.1
+rsa==4.9
+scikit-learn==1.2.2
+scipy==1.10.1
+six==1.16.0
+smt==1.3.0
+sympy==1.12
+tensorboard==2.13.0
+tensorboard-data-server==0.7.0
+tensorboard-plugin-profile==2.13.0
+tensorstore==0.1.37
+threadpoolctl==3.1.0
+toolz==0.12.0
+torch==2.0.1
+torch-tb-profiler==0.4.1
+torchmetrics==0.11.4
+tqdm==4.65.0
+triton==2.0.0
+typing-extensions==4.6.2
+tzdata==2023.3
+urllib3==1.26.16
+virtualenv==20.16.5
+Werkzeug==2.3.4
+yarl==1.9.2