Traceback (most recent call last):
  File "/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/site-packages/jupyter_cache/executors/utils.py", line 51, in single_nb_execution
    executenb(
  File "/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/site-packages/nbclient/client.py", line 1204, in execute
    return NotebookClient(nb=nb, resources=resources, km=km, **kwargs).execute()
  File "/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/site-packages/nbclient/util.py", line 84, in wrapped
    return just_run(coro(*args, **kwargs))
  File "/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/site-packages/nbclient/util.py", line 62, in just_run
    return loop.run_until_complete(coro)
  File "/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/asyncio/base_events.py", line 642, in run_until_complete
    return future.result()
  File "/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/site-packages/nbclient/client.py", line 663, in async_execute
    await self.async_execute_cell(
  File "/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/site-packages/nbclient/client.py", line 965, in async_execute_cell
    await self._check_raise_for_error(cell, cell_index, exec_reply)
  File "/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/site-packages/nbclient/client.py", line 862, in _check_raise_for_error
    raise CellExecutionError.from_cell_and_msg(cell, exec_reply_content)
nbclient.exceptions.CellExecutionError: An error occurred while executing the following cell:
------------------
%matplotlib inline

import time
import numpy as np
import tensorflow as tf
from matplotlib import image
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from IPython.display import display

np.random.seed(2021)
------------------

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [1], in <cell line: 5>()
      3 import time
      4 import numpy as np
----> 5 import tensorflow as tf
      6 from matplotlib import image
      7 import matplotlib.pyplot as plt

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/__init__.py:51, in <module>
     49 from ._api.v2 import autograph
     50 from ._api.v2 import bitwise
---> 51 from ._api.v2 import compat
     52 from ._api.v2 import config
     53 from ._api.v2 import data

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/__init__.py:37, in <module>
      3 """Compatibility functions.
      4 
      5 The `tf.compat` module contains two sets of compatibility functions.
   (...)
     32 
     33 """
     35 import sys as _sys
---> 37 from . import v1
     38 from . import v2
     39 from tensorflow.python.compat.compat import forward_compatibility_horizon

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/__init__.py:30, in <module>
     28 from . import autograph
     29 from . import bitwise
---> 30 from . import compat
     31 from . import config
     32 from . import data

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/__init__.py:37, in <module>
      3 """Compatibility functions.
      4 
      5 The `tf.compat` module contains two sets of compatibility functions.
   (...)
     32 
     33 """
     35 import sys as _sys
---> 37 from . import v1
     38 from . import v2
     39 from tensorflow.python.compat.compat import forward_compatibility_horizon

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/v1/__init__.py:47, in <module>
     45 from tensorflow._api.v2.compat.v1 import layers
     46 from tensorflow._api.v2.compat.v1 import linalg
---> 47 from tensorflow._api.v2.compat.v1 import lite
     48 from tensorflow._api.v2.compat.v1 import logging
     49 from tensorflow._api.v2.compat.v1 import lookup

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/__init__.py:9, in <module>
      6 import sys as _sys
      8 from . import constants
----> 9 from . import experimental
     10 from tensorflow.lite.python.lite import Interpreter
     11 from tensorflow.lite.python.lite import OpHint

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/__init__.py:8, in <module>
      3 """Public API for tf.lite.experimental namespace.
      4 """
      6 import sys as _sys
----> 8 from . import authoring
      9 from tensorflow.lite.python.analyzer import ModelAnalyzer as Analyzer
     10 from tensorflow.lite.python.lite import OpResolverType

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/authoring/__init__.py:8, in <module>
      3 """Public API for tf.lite.experimental.authoring namespace.
      4 """
      6 import sys as _sys
----> 8 from tensorflow.lite.python.authoring.authoring import compatible

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/lite/python/authoring/authoring.py:43, in <module>
     39 import functools
     42 # pylint: disable=g-import-not-at-top
---> 43 from tensorflow.lite.python import convert
     44 from tensorflow.lite.python import lite
     45 from tensorflow.lite.python.metrics import converter_error_data_pb2

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/lite/python/convert.py:29, in <module>
     26 import six
     28 from tensorflow.lite.python import lite_constants
---> 29 from tensorflow.lite.python import util
     30 from tensorflow.lite.python import wrap_toco
     31 from tensorflow.lite.python.convert_phase import Component

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/tensorflow/lite/python/util.py:51, in <module>
     47 # Jax functions used by TFLite
     48 # pylint: disable=g-import-not-at-top
     49 # pylint: disable=unused-import
     50 try:
---> 51   from jax import xla_computation as _xla_computation
     52 except ImportError:
     53   _xla_computation = None

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/__init__.py:116, in <module>
     40 from ._src.config import (
     41   config as config,
     42   enable_checks as enable_checks,
   (...)
     51   numpy_rank_promotion as numpy_rank_promotion,
     52 )
     53 from ._src.api import (
     54   ad,  # TODO(phawkins): update users to avoid this.
     55   checkpoint as checkpoint,
   (...)
    114   xla_computation as xla_computation,
    115 )
--> 116 from .experimental.maps import soft_pmap as soft_pmap
    117 from .version import __version__ as __version__
    119 # These submodules are separate because they are in an import cycle with
    120 # jax and rely on the names imported above.

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/experimental/maps.py:26, in <module>
     23 from functools import wraps, partial, partialmethod
     24 from enum import Enum
---> 26 from .. import numpy as jnp
     27 from .. import core
     28 from .. import linear_util as lu

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/numpy/__init__.py:19, in <module>
      1 # Copyright 2018 Google LLC
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     17 
     18 # flake8: noqa: F401
---> 19 from . import fft as fft
     20 from . import linalg as linalg
     22 from jax.interpreters.xla import DeviceArray as DeviceArray

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/numpy/fft.py:17, in <module>
      1 # Copyright 2020 Google LLC
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     14 
     15 # flake8: noqa: F401
---> 17 from jax._src.numpy.fft import (
     18   ifft as ifft,
     19   ifft2 as ifft2,
     20   ifftn as ifftn,
     21   ifftshift as ifftshift,
     22   ihfft as ihfft,
     23   irfft as irfft,
     24   irfft2 as irfft2,
     25   irfftn as irfftn,
     26   fft as fft,
     27   fft2 as fft2,
     28   fftfreq as fftfreq,
     29   fftn as fftn,
     30   fftshift as fftshift,
     31   hfft as hfft,
     32   rfft as rfft,
     33   rfft2 as rfft2,
     34   rfftfreq as rfftfreq,
     35   rfftn as rfftn,
     36 )
     38 # Module initialization is encapsulated in a function to avoid accidental
     39 # namespace pollution.
     40 _NOT_IMPLEMENTED = []

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/_src/numpy/fft.py:19, in <module>
     16 import operator
     17 import numpy as np
---> 19 from jax import lax
     20 from jax._src.lib import xla_client
     21 from jax._src.util import safe_zip

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/lax/__init__.py:332, in <module>
    299 from jax._src.lax.lax import (_reduce_sum, _reduce_max, _reduce_min, _reduce_or,
    300                   _reduce_and, _reduce_window_sum, _reduce_window_max,
    301                   _reduce_window_min, _reduce_window_prod,
   (...)
    306                   _upcast_fp16_for_computation, _broadcasting_shape_rule,
    307                   _eye, _tri, _delta, _ones, _zeros, _dilate_shape)
    308 from jax._src.lax.control_flow import (
    309   associative_scan as associative_scan,
    310   cond as cond,
   (...)
    330   while_p as while_p,
    331 )
--> 332 from jax._src.lax.fft import (
    333   fft as fft,
    334   fft_p as fft_p,
    335 )
    336 from jax._src.lax.parallel import (
    337   all_gather as all_gather,
    338   all_to_all as all_to_all,
   (...)
    355   xeinsum as xeinsum,
    356 )
    357 from jax._src.lax.other import (
    358   conv_general_dilated_patches as conv_general_dilated_patches
    359 )

File ~/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/_src/lax/fft.py:145, in <module>
    143 batching.primitive_batchers[fft_p] = fft_batching_rule
    144 if pocketfft:
--> 145   xla.backend_specific_translations['cpu'][fft_p] = pocketfft.pocketfft

AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft'
AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft'