To import these modules from anywhere, run pip install -e .
To install all dependencies, do pip install -r requirements.txt
To see the model in action, start at example.ipynb, example2.ipynb or experiments/supervised_example/. example2.ipynb shows an example of writing a custom task, NestedDMS. This task was designed with Matt Rosen and Nabilah Sammudin.
Note that the parameters in xdg_parameters.py affect only on the data generator xdg_stimulus.py, and not the model described. RNN parameters sitting in xdg_paramters.py are a remnant of how this file was used in a previous repo. Same for stp_parameters.py and stp_stimulus.py.
For model description, see the Methods section of Circuit mechanisms for the maintenance and manipulation of information in working memory.
The model was written as a Sonnet module so that you can compose it with other Sonnet modules. I recommend working with Sonnet modules because of their simplicity and readable source code. If you want to use Keras modules, note that Sonnet modules are not always composable with Keras modules. Keras functions are okay though (like tf.keras.optimizers and tf.keras.initializers). If you want to use Keras modules, you'd probably want to rewrite the model as a Keras module. This should be easy enough, you'll just have to rename a few things (see this and this).
If you're new to TensorFlow 2, I wrote this tutorial notebook.