Can't use TFDS on M1 MAC

/!\ PLEASE INCLUDE THE FULL STACKTRACE AND CODE SNIPPET

Short description
Hello,
I hope you're doing great.
I am using TFDS on a M1 MAC and try to load LVIS.

Environment information

• Operating System: OSX

• Python version: 3.10
  tensorflow-addons==0.19.0
  tensorflow-datasets==4.8.2
  tensorflow-estimator==2.11.0
  tensorflow-macos==2.11.0
  tensorflow-metadata==1.12.0
  tensorflow-metal==0.7.0

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

Reproduction instructions

tfds.load('lvis')

If you share a colab, make sure to update the permissions to share it.

Link to logs
Traceback (most recent call last):
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/lazy_imports_lib.py", line 30, in _try_import
mod = importlib.import_module(module_name)
File "/Users/colombo/miniforge3/lib/python3.10/importlib/init.py", line 126, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
File "", line 1050, in _gcd_import
File "", line 1027, in _find_and_load
File "", line 1006, in _find_and_load_unlocked
File "", line 688, in _load_unlocked
File "", line 883, in exec_module
File "", line 241, in _call_with_frames_removed
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/apache_beam/init.py", line 92, in
from apache_beam import coders
File "/Applications/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_bundle/pydev_import_hook.py", line 21, in do_import
module = self._system_import(name, *args, **kwargs)
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/apache_beam/coders/init.py", line 17, in
from apache_beam.coders.coders import *
File "/Applications/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_bundle/pydev_import_hook.py", line 21, in do_import
module = self._system_import(name, *args, **kwargs)
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/apache_beam/coders/coders.py", line 59, in
from apache_beam.coders import coder_impl
File "/Applications/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_bundle/pydev_import_hook.py", line 21, in do_import
module = self._system_import(name, *args, **kwargs)
ImportError: dlopen(/Users/colombo/miniforge3/lib/python3.10/site-packages/apache_beam/coders/coder_impl.cpython-310-darwin.so, 0x0002): tried: '/Users/colombo/miniforge3/lib/python3.10/site-packages/apache_beam/coders/coder_impl.cpython-310-darwin.so' (mach-o file, but is an incompatible architecture (have 'x86_64', need 'arm64e'))
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/Users/colombo/miniforge3/lib/python3.10/code.py", line 90, in runcode
exec(code, self.locals)
File "", line 1, in
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/logging/init.py", line 169, in call
return function(*args, **kwargs)
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/load.py", line 617, in load
dbuilder.download_and_prepare(**download_and_prepare_kwargs)
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/logging/init.py", line 169, in call
return function(*args, **kwargs)
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/dataset_builder.py", line 628, in download_and_prepare
self._download_and_prepare(
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/dataset_builder.py", line 1436, in _download_and_prepare
split_generators = self._split_generators( # pylint: disable=unexpected-keyword-arg
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/datasets/lvis/lvis_dataset_builder.py", line 110, in _split_generators
tfds.Split.TRAIN: self._generate_examples(
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/datasets/lvis/lvis_dataset_builder.py", line 157, in _generate_examples
beam = tfds.core.lazy_imports.apache_beam
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/utils/py_utils.py", line 154, in get
return self.fget.get(None, objtype)() # pytype: disable=attribute-error
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/lazy_imports_lib.py", line 53, in apache_beam
return _try_import("apache_beam")
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/lazy_imports_lib.py", line 39, in _try_import
utils.reraise(e, suffix=err_msg)
File "/Users/colombo/miniforge3/lib/python3.10/site-packages/tensorflow_datasets/core/utils/py_utils.py", line 401, in reraise
raise exception from e
ImportError: dlopen(/Users/colombo/miniforge3/lib/python3.10/site-packages/apache_beam/coders/coder_impl.cpython-310-darwin.so, 0x0002): tried: '/Users/colombo/miniforge3/lib/python3.10/site-packages/apache_beam/coders/coder_impl.cpython-310-darwin.so' (mach-o file, but is an incompatible architecture (have 'x86_64', need 'arm64e'))
Failed importing apache_beam. This likely means that the dataset requires additional dependencies that have to be manually installed (usually with pip install apache_beam). See setup.py extras_require.

Expected behavior
Beeing abloe to load LVIS.

Additional context
Add any other context about the problem here.

Reproduce tfds.load('lvis') on the reported M1/macOS environment, then inspect tensorflow_datasets/datasets/lvis/lvis_dataset_builder.py and core/lazy_imports_lib.py around the Apache Beam import. Check the reported x86_64-versus-arm64e failure and the project's Apache Beam dependency guidance; done means the issue has a verified compatible path or a documented, tested project-level fix.