Problem importing TensorFlow probability in TFX pipeline on Google Cloud AI Platform

Hello,

I have a TFX pipeline with a model using TensorFlow Probability deployed using the Kubeflow engine on the Google Cloud AI Platform.
The cloud TensorFlow version is 2.4.
However, when importing TensorFlow Probability I get the following error, thus blocking the pipeline:

Traceback (most recent call last):
File "/opt/conda/lib/python3.7/runpy.py", line 193, in _run_module_as_main
"main", mod_spec)
File "/opt/conda/lib/python3.7/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/opt/conda/lib/python3.7/site-packages/tfx/orchestration/kubeflow/container_entrypoint.py", line 364, in
main()
File "/opt/conda/lib/python3.7/site-packages/tfx/orchestration/kubeflow/container_entrypoint.py", line 357, in main
execution_info = launcher.launch()
File "/opt/conda/lib/python3.7/site-packages/tfx/orchestration/launcher/base_component_launcher.py", line 209, in launch
copy.deepcopy(execution_decision.exec_properties))
File "/opt/conda/lib/python3.7/site-packages/tfx/orchestration/launcher/in_process_component_launcher.py", line 72, in _run_executor
copy.deepcopy(input_dict), output_dict, copy.deepcopy(exec_properties))
File "/opt/conda/lib/python3.7/site-packages/tfx/components/trainer/executor.py", line 182, in Do
run_fn = udf_utils.get_fn(exec_properties, 'run_fn')
File "/opt/conda/lib/python3.7/site-packages/tfx/components/util/udf_utils.py", line 49, in get_fn
exec_properties[_MODULE_FILE_KEY], fn_name)
File "/opt/conda/lib/python3.7/site-packages/tfx/utils/import_utils.py", line 127, in import_func_from_source
loader.exec_module(module)
File "", line 728, in exec_module
File "", line 219, in _call_with_frames_removed
File "./censeo/models/vae_trainer.py", line 10, in
import tensorflow_probability as tfp
File "/opt/conda/lib/python3.7/site-packages/tensorflow_probability/init.py", line 77, in
from tensorflow_probability.python import * # pylint: disable=wildcard-import
File "/opt/conda/lib/python3.7/site-packages/tensorflow_probability/python/init.py", line 23, in
from tensorflow_probability.python import distributions
File "/opt/conda/lib/python3.7/site-packages/tensorflow_probability/python/distributions/init.py", line 94, in
from tensorflow_probability.python.distributions.pixel_cnn import PixelCNN
File "/opt/conda/lib/python3.7/site-packages/tensorflow_probability/python/distributions/pixel_cnn.py", line 37, in
from tensorflow_probability.python.layers import weight_norm
File "/opt/conda/lib/python3.7/site-packages/tensorflow_probability/python/layers/init.py", line 31, in
from tensorflow_probability.python.layers.distribution_layer import CategoricalMixtureOfOneHotCategorical
File "/opt/conda/lib/python3.7/site-packages/tensorflow_probability/python/layers/distribution_layer.py", line 28, in
from cloudpickle.cloudpickle import CloudPickler
ImportError: cannot import name 'CloudPickler' from 'cloudpickle.cloudpickle' (/opt/conda/lib/python3.7/site-packages/cloudpickle/cloudpickle.py)

Do you have an idea to fix the problem? The version available in the Kubeflow cloud instance seems broken.

Thanks for your help

Start with the failing import in censeo/models/vae_trainer.py and trace the TensorFlow Probability imports through distribution_layer.py and cloudpickle.cloudpickle. Reproduce the import in the stated Python 3.7, TensorFlow 2.4, TFX, and Kubeflow environment; done means TensorFlow Probability imports successfully in the Google Cloud AI Platform pipeline.