Tensorflow on Apple M1 cannot import

Hi,

OS: macOS 12.4 CPU: Apple M1

I cannot import the new TensorFlow 2.9.0 on Apple M1. I got an error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/Users/martin/Documents/Projects/Solar-Transformer/Testing.ipynb Cell 3' in <cell line: 1>()
----> 1 from tensorflow.keras.layers import Add, Dense, Dropout, Layer, LayerNormalization, MultiHeadAttention, Normalization
      2 from tensorflow.keras.models import Model
      3 from tensorflow.keras.initializers import TruncatedNormal

File ~/miniforge3/lib/python3.9/site-packages/tensorflow/__init__.py:37, in <module>
     34 import sys as _sys
     35 import typing as _typing
---> 37 from tensorflow.python.tools import module_util as _module_util
     38 from tensorflow.python.util.lazy_loader import LazyLoader as _LazyLoader
     40 # Make sure code inside the TensorFlow codebase can use tf2.enabled() at import.

File ~/miniforge3/lib/python3.9/site-packages/tensorflow/python/__init__.py:42, in <module>
     37 from tensorflow.python.eager import context
     39 # pylint: enable=wildcard-import
     40 
     41 # Bring in subpackages.
---> 42 from tensorflow.python import data
     43 from tensorflow.python import distribute
     44 # from tensorflow.python import keras

File ~/miniforge3/lib/python3.9/site-packages/tensorflow/python/data/__init__.py:21, in <module>
     15 """`tf.data.Dataset` API for input pipelines.
     16 
     17 See [Importing Data](https://tensorflow.org/guide/data) for an overview.
     18 """
     20 # pylint: disable=unused-import
---> 21 from tensorflow.python.data import experimental
     22 from tensorflow.python.data.ops.dataset_ops import AUTOTUNE
     23 from tensorflow.python.data.ops.dataset_ops import Dataset

File ~/miniforge3/lib/python3.9/site-packages/tensorflow/python/data/util/structure.py:22, in <module>
     19 import six
     20 import wrapt
---> 22 from tensorflow.python.data.util import nest
     23 from tensorflow.python.framework import composite_tensor
     24 from tensorflow.python.framework import ops

File ~/miniforge3/lib/python3.9/site-packages/tensorflow/python/data/util/nest.py:36, in <module>
     16 """## Functions for working with arbitrarily nested sequences of elements.
     17 
     18 NOTE(mrry): This fork of the `tensorflow.python.util.nest` module
 (...)
     31    arrays.
     32 """
     34 import six as _six
---> 36 from tensorflow.python.framework import sparse_tensor as _sparse_tensor
     37 from tensorflow.python.util import _pywrap_utils
     38 from tensorflow.python.util import nest

File ~/miniforge3/lib/python3.9/site-packages/tensorflow/python/framework/sparse_tensor.py:24, in <module>
     22 from tensorflow.python import tf2
     23 from tensorflow.python.framework import composite_tensor
---> 24 from tensorflow.python.framework import constant_op
     25 from tensorflow.python.framework import dtypes
     26 from tensorflow.python.framework import ops

File ~/miniforge3/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py:25, in <module>
     23 from tensorflow.core.framework import types_pb2
     24 from tensorflow.python.eager import context
---> 25 from tensorflow.python.eager import execute
     26 from tensorflow.python.framework import dtypes
     27 from tensorflow.python.framework import op_callbacks

File ~/miniforge3/lib/python3.9/site-packages/tensorflow/python/eager/execute.py:23, in <module>
     21 from tensorflow.python import pywrap_tfe
     22 from tensorflow.python.eager import core
---> 23 from tensorflow.python.framework import dtypes
     24 from tensorflow.python.framework import ops
     25 from tensorflow.python.framework import tensor_shape

File ~/miniforge3/lib/python3.9/site-packages/tensorflow/python/framework/dtypes.py:29, in <module>
     26 from tensorflow.python.lib.core import _pywrap_bfloat16
     27 from tensorflow.python.util.tf_export import tf_export
---> 29 _np_bfloat16 = _pywrap_bfloat16.TF_bfloat16_type()
     32 @tf_export("dtypes.DType", "DType")
     33 class DType(_dtypes.DType):
     34   """Represents the type of the elements in a `Tensor`.
     35 
     36   `DType`'s are used to specify the output data type for operations which
   (...)
     46   See `tf.dtypes` for a complete list of `DType`'s defined.
     47   """

Example code:

from tensorflow.keras.layers import Add, Dense, Dropout, Layer, LayerNormalization, MultiHeadAttention, Normalization
from tensorflow.keras.models import Model
from tensorflow.keras.initializers import TruncatedNormal
from tensorflow.keras.utils import timeseries_dataset_from_array

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np