NotImplementedError in Ubuntu 24.04 (WSL2) with GPU

[MatteoOlarte]

The tfds.load('plant_leaves', with_info=True) function works perfectly on Windows 11, but when running the same code on Ubuntu 24.04 (WSL2), the following error appears:

File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/reader.py", line 70, in _get_dataset_from_filename
    ds = file_adapters.ADAPTER_FOR_FORMAT[file_format].make_tf_data(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/file_adapters.py", line 301, in make_tf_data
    raise NotImplementedError(

NotImplementedError: `.as_dataset()` not implemented for ArrayRecord files. Please, use `.as_data_source()`.

Environment information

• Operating System: Ubuntu 24.04 (WSL2)

• Python version: 3.12.8

• tensorflow-datasets/tfds-nightly version: 4.9.7.dev202502260044

• tensorflow/tf-nightly version: 2.18.0

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

Reproduction instructions

import tensorflow as tf
import tensorflow_datasets as tf_datasets


dataset, metadata = tf_datasets.load('plant_leaves', as_supervised=True, with_info=True)

Link to logs

Traceback (most recent call last):
  File "/mnt/d/Users/mateo/Documents/Python Development/NeuronalNetworks/main.py", line 10, in <module>
    tf_datasets.load('plant_leaves', with_info=True, as_supervised=True)
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/logging/__init__.py", line 176, in __call__
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/load.py", line 681, in load
    ds = dbuilder.as_dataset(**as_dataset_kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/logging/__init__.py", line 176, in __call__
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/dataset_builder.py", line 1046, in as_dataset
    all_ds = tree.map_structure(build_single_dataset, split)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tree/__init__.py", line 428, in map_structure
    [func(*args) for args in zip(*map(flatten, structures))])
     ^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/dataset_builder.py", line 1064, in _build_single_dataset
    ds = self._as_dataset(
         ^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/dataset_builder.py", line 1553, in _as_dataset
    return reader.read(
           ^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/reader.py", line 433, in read
    return tree.map_structure(_read_instruction_to_ds, instructions)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tree/__init__.py", line 428, in map_structure
    [func(*args) for args in zip(*map(flatten, structures))])
     ^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/reader.py", line 425, in _read_instruction_to_ds
    return self.read_files(
           ^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/reader.py", line 465, in read_files
    ds = _read_files(
         ^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/reader.py", line 303, in _read_files
    ds = instruction_ds.interleave(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 2534, in interleave
    return interleave_op._interleave(self, map_func, cycle_length, block_length,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/data/ops/interleave_op.py", line 49, in _interleave
    return _ParallelInterleaveDataset(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/data/ops/interleave_op.py", line 119, in __init__
    self._map_func = structured_function.StructuredFunctionWrapper(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/data/ops/structured_function.py", line 265, in __init__
    self._function = fn_factory()
                     ^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1251, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1221, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 696, in _initialize
    self._concrete_variable_creation_fn = tracing_compilation.trace_function(
                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 178, in trace_function
    concrete_function = _maybe_define_function(
                        ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 283, in _maybe_define_function
    concrete_function = _create_concrete_function(
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 310, in _create_concrete_function
    traced_func_graph = func_graph_module.func_graph_from_py_func(
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/framework/func_graph.py", line 1059, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 599, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/data/ops/structured_function.py", line 231, in wrapped_fn
    ret = wrapper_helper(*args)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/data/ops/structured_function.py", line 161, in wrapper_helper
    ret = autograph.tf_convert(self._func, ag_ctx)(*nested_args)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/autograph/impl/api.py", line 690, in wrapper
    return converted_call(f, args, kwargs, options=options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/autograph/impl/api.py", line 352, in converted_call
    return converted_call(
           ^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/autograph/impl/api.py", line 377, in converted_call
    return _call_unconverted(f, args, kwargs, options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow/python/autograph/impl/api.py", line 459, in _call_unconverted
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/reader.py", line 70, in _get_dataset_from_filename
    ds = file_adapters.ADAPTER_FOR_FORMAT[file_format].make_tf_data(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/anaconda3/envs/tensorflow-cuda/lib/python3.12/site-packages/tensorflow_datasets/core/file_adapters.py", line 301, in make_tf_data
    raise NotImplementedError(
NotImplementedError: `.as_dataset()` not implemented for ArrayRecord files. Please, use `.as_data_source()`.

Expected behavior
The dataset should load correctly on Ubuntu 24.04 (WSL2) with GPU, just as it does on Windows 11.

Additional context

• The error only appears when switching from Windows 11 to Ubuntu 24.04 (WSL2).
• The dataset does load correctly in Windows, but fails in Linux.