Not compatible with `tensorflow-macos`

Importing distrax causes the following stack trace error

```
...
import distrax
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/distrax/__init__.py", line 18, in
from distrax._src.bijectors.bijector import Bijector
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/distrax/_src/bijectors/bijector.py", line 26, in
tfb = tfp.bijectors
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 53, in __getattr__
module = self._load()
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 40, in _load
module = importlib.import_module(self.__name__)
File "/Users/akbir.khan/.pyenv/versions/3.9.12/lib/python3.9/importlib/__init__.py", line 127, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/__init__.py", line 41, in
from tensorflow_probability.substrates.jax import bijectors
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/bijectors/__init__.py", line 19, in
from tensorflow_probability.substrates.jax.bijectors.absolute_value import AbsoluteValue
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/bijectors/absolute_value.py", line 17, in
from tensorflow_probability.python.internal.backend.jax.compat import v2 as tf
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow_probability/python/internal/backend/jax/__init__.py", line 18, in
from tensorflow_probability.python.internal.backend.jax import bitwise
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow_probability/python/internal/backend/jax/bitwise.py", line 19, in
from tensorflow_probability.python.internal.backend.jax import _utils as utils
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow_probability/python/internal/backend/jax/_utils.py", line 25, in
from tensorflow.python.ops import array_ops # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top,unused-import
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/__init__.py", line 37, in
from tensorflow.python.tools import module_util as _module_util
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/__init__.py", line 42, in
from tensorflow.python import data
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/data/__init__.py", line 21, in
from tensorflow.python.data import experimental
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/data/experimental/__init__.py", line 96, in
from tensorflow.python.data.experimental import service
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/data/experimental/service/__init__.py", line 419, in
from tensorflow.python.data.experimental.ops.data_service_ops import distribute
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/data/experimental/ops/data_service_ops.py", line 22, in
from tensorflow.python.data.experimental.ops import compression_ops
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/data/experimental/ops/compression_ops.py", line 16, in
from tensorflow.python.data.util import structure
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/data/util/structure.py", line 22, in
from tensorflow.python.data.util import nest
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/data/util/nest.py", line 34, in
from tensorflow.python.framework import sparse_tensor as _sparse_tensor
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/framework/sparse_tensor.py", line 24, in
from tensorflow.python.framework import constant_op
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py", line 25, in
from tensorflow.python.eager import execute
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/eager/execute.py", line 21, in
from tensorflow.python.framework import dtypes
File "/Users/akbir.khan/pax/venv/lib/python3.9/site-packages/tensorflow/python/framework/dtypes.py", line 34, in
_np_bfloat16 = _pywrap_bfloat16.TF_bfloat16_type()
TypeError: Unable to convert function return value to a Python type! The signature was
() -> handle
```