AttributeError: module 'tree' has no attribute 'map_structure'

[rafadotnechi]

/!\ PLEASE INCLUDE THE FULL STACKTRACE AND CODE SNIPPET

Short description
Description of the bug.

Environment information

• Operating System: Windows 10

• Python version: Python 3.12.4

• tensorflow-datasets/tfds-nightly tensorflow-datasets / version: Version: 4.9.6

• tensorflow/tf-nightly version: tensorflow / Version: 2.16.1

• Does the issue still exists with the last tfds-nightly package (pip install --upgrade tfds-nightly) ?

Yes

Reproduction instructions
load any dataset

import tensorflow_datasets as tfds
import tensorflow as tf

# Construct a tf.data.Dataset
ds = tfds.load('mnist', split='train', as_supervised=True, shuffle_files=True)

Stacktrace

Traceback (most recent call last):
  File "C:/Users/rafadotnechi/Documents/PythonDev/load_test.py", line 5, in <module>
    ds = tfds.load('mnist', split='train', as_supervised=True, shuffle_files=True)
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\core\logging\__init__.py", line 176, in __call__
    return function(*args, **kwargs)
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\core\load.py", line 647, in load
    _download_and_prepare_builder(dbuilder, download, download_and_prepare_kwargs)
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\core\load.py", line 506, in _download_and_prepare_builder
    dbuilder.download_and_prepare(**download_and_prepare_kwargs)
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\core\logging\__init__.py", line 176, in __call__
    return function(*args, **kwargs)
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\core\dataset_builder.py", line 699, in download_and_prepare
    self._download_and_prepare(
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\core\dataset_builder.py", line 1669, in _download_and_prepare
    split_infos = self._generate_splits(dl_manager, download_config)
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\core\dataset_builder.py", line 1620, in _generate_splits
    split_generators = self._split_generators(  # pylint: disable=unexpected-keyword-arg
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\image_classification\mnist.py", line 119, in _split_generators
    mnist_files = dl_manager.download_and_extract(
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\core\download\download_manager.py", line 694, in download_and_extract
    return _map_promise(self._download_extract, url_or_urls)
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\tensorflow_datasets\core\download\download_manager.py", line 790, in _map_promise
    all_promises = tree.map_structure(map_fn, all_inputs)  # Apply the function
  File "C:\Users\rafadotnechi\AppData\Local\Programs\Python\Python312\Lib\site-packages\etils\epy\lazy_imports_utils.py", line 109, in __getattr__
    return getattr(self._module, name)
AttributeError: module 'tree' has no attribute 'map_structure'