Fix import statements in some configs

PiperOrigin-RevId: 681846808 Change-Id: Id0b574ebc3d0d8e943f1d4399b6cbc04382dbfe3
google-deepmind · Oct 3, 2024 · 38df5df · 38df5df
1 parent fa11e80
commit 38df5df
Show file tree

Hide file tree

Showing 7 changed files with 14 additions and 14 deletions.
diff --git a/ferminet/configs/excited/atoms.py b/ferminet/configs/excited/atoms.py
@@ -15,8 +15,8 @@
 """Config to reproduce Fig. 1 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs import atom
-from ferminet.jax.configs.excited import presets
+from ferminet.configs import atom
+from ferminet.configs.excited import presets
 import ml_collections
 
 

diff --git a/ferminet/configs/excited/benzene.py b/ferminet/configs/excited/benzene.py
@@ -15,7 +15,7 @@
 """Config to reproduce Fig. 6 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 

diff --git a/ferminet/configs/excited/carbon_dimer.py b/ferminet/configs/excited/carbon_dimer.py
@@ -15,7 +15,7 @@
 """Config to reproduce Fig. 3 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 

diff --git a/ferminet/configs/excited/double_excitation.py b/ferminet/configs/excited/double_excitation.py
@@ -15,7 +15,7 @@
 """Config to reproduce Fig. 5 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 

diff --git a/ferminet/configs/excited/oscillator.py b/ferminet/configs/excited/oscillator.py
@@ -16,7 +16,7 @@
 
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 import pyscf

diff --git a/ferminet/configs/excited/twisted_ethylene.py b/ferminet/configs/excited/twisted_ethylene.py
@@ -15,7 +15,7 @@
 """Config to reproduce Fig. 4 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 import numpy as np

diff --git a/ferminet/tests/excited_test.py b/ferminet/tests/excited_test.py
@@ -23,13 +23,13 @@
 import chex
 from ferminet import base_config
 from ferminet import train
-from ferminet.jax.configs.excited import atoms
-from ferminet.jax.configs.excited import benzene
-from ferminet.jax.configs.excited import carbon_dimer
-from ferminet.jax.configs.excited import double_excitation
-from ferminet.jax.configs.excited import oscillator
-from ferminet.jax.configs.excited import presets
-from ferminet.jax.configs.excited import twisted_ethylene
+from ferminet.configs.excited import atoms
+from ferminet.configs.excited import benzene
+from ferminet.configs.excited import carbon_dimer
+from ferminet.configs.excited import double_excitation
+from ferminet.configs.excited import oscillator
+from ferminet.configs.excited import presets
+from ferminet.configs.excited import twisted_ethylene
 import jax
 
 import pyscf