Data download issue with official Flax Image Net example

Question

I am still trying to understand this official Flax Example. For the convenience of the experiment, I have created my own copy.

In the section on running locally, it seems that there is no download command. Therefore when I run python main.py --workdir=./imagenet --config=configs/v100_x8.py, I got the error:

Traceback (most recent call last):
  File "/content/FlaxImageNet/main.py", line 65, in <module>
    app.run(main)
  File "/usr/local/lib/python3.10/dist-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/usr/local/lib/python3.10/dist-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "/content/FlaxImageNet/main.py", line 60, in main
    train.train_and_evaluate(FLAGS.config, FLAGS.workdir)
  File "/content/FlaxImageNet/train.py", line 280, in train_and_evaluate
    train_iter = create_input_iter(
  File "/content/FlaxImageNet/train.py", line 185, in create_input_iter
    ds = input_pipeline.create_split(
  File "/content/FlaxImageNet/input_pipeline.py", line 207, in create_split
    ds = dataset_builder.as_dataset(split=split, decoders={
  File "/usr/local/lib/python3.10/dist-packages/tensorflow_datasets/core/logging/__init__.py", line 81, in decorator
    return function(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/tensorflow_datasets/core/dataset_builder.py", line 542, in as_dataset
    raise AssertionError(
AssertionError: Dataset imagenet2012: could not find data in /root/tensorflow_datasets. Please make sure to call dataset_builder.download_and_prepare(), or pass download=True to tfds.load() before trying to access the tf.data.Dataset object.

Now when I call the download command later, using the command

python -c "
import tensorflow_datasets as tfds
tfds.builder('imagenet2012').download_and_prepare(
    download_config=tfds.download.DownloadConfig(
        manual_dir="/root/tensorflow_datasets"))
"

It reports a strange error,

File "<string>", line 5
    manual_dir=/root/tensorflow_datasets))
               ^
SyntaxError: invalid syntax

It is strange since I have typed in the quotation marks, and in the console of Colab, it seems just to ignore it.

For your information. I am using Jax and Jaxlib==0.4.6