TensorFlow v2.14.0 breaks TensorFlow Probability at import

Summary of problem

Today's release of TensorFlow v2.14.0 breaks TensorFlow Probability at import.

In a fresh Python 3.11 virtual environment, installation of tensorflow v2.14.0 and tensorflow-probability v0.21.0 causes a

ValueError: Arg specs do not match: original=FullArgSpec(args=['input', 'dtype', 'name', 'layout'], varargs=None, varkw=None, defaults=(None, None, None), kwonlyargs=[], kwonlydefaults=None, annotations={}), new=FullArgSpec(args=['input', 'dtype', 'name'], varargs=None, varkw=None, defaults=(None, None), kwonlyargs=[], kwonlydefaults=None, annotations={}), fn=<function ones_like_v2 at 0x7fce4da0e480>

at import of both.

Reproducible example

$ docker run --rm -ti python:3.11 /bin/bash 
root@c4673c623033:/# python -m venv venv && . venv/bin/activate
(venv) root@c4673c623033:/# python -m pip --quiet install --upgrade pip setuptools wheel
(venv) root@c4673c623033:/# python -m pip --quiet install --upgrade 'tensorflow==2.14.0' 'tensorflow-probability==0.21.0'
(venv) root@c4673c623033:/# python -m pip list | grep tensorflow
tensorflow                   2.14.0
tensorflow-estimator         2.14.0
tensorflow-io-gcs-filesystem 0.34.0
tensorflow-probability       0.21.0
(venv) root@c4673c623033:/# python -c 'import tensorflow; import tensorflow_probability'
2023-09-26 23:27:07.453221: I tensorflow/core/util/port.cc:111] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-09-26 23:27:07.454646: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.
2023-09-26 23:27:07.478170: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-09-26 23:27:07.478209: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-09-26 23:27:07.478229: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2023-09-26 23:27:07.483495: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.
2023-09-26 23:27:07.483705: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-09-26 23:27:08.160315: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/__init__.py", line 20, in <module>
    from tensorflow_probability import substrates
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/substrates/__init__.py", line 17, in <module>
    from tensorflow_probability.python.internal import all_util
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/__init__.py", line 138, in <module>
    dir(globals()[pkg_name])  # Forces loading the package from its lazy loader.
    ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 57, in __dir__
    module = self._load()
             ^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 40, in _load
    module = importlib.import_module(self.__name__)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/experimental/__init__.py", line 31, in <module>
    from tensorflow_probability.python.experimental import bayesopt
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/experimental/bayesopt/__init__.py", line 17, in <module>
    from tensorflow_probability.python.experimental.bayesopt import acquisition
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/experimental/bayesopt/acquisition/__init__.py", line 17, in <module>
    from tensorflow_probability.python.experimental.bayesopt.acquisition.acquisition_function import AcquisitionFunction
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/experimental/bayesopt/acquisition/acquisition_function.py", line 22, in <module>
    from tensorflow_probability.python.internal import prefer_static as ps
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/internal/prefer_static.py", line 361, in <module>
    ones_like = _copy_docstring(tf.ones_like, _ones_like)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/internal/prefer_static.py", line 84, in _copy_docstring
    raise ValueError(
ValueError: Arg specs do not match: original=FullArgSpec(args=['input', 'dtype', 'name', 'layout'], varargs=None, varkw=None, defaults=(None, None, None), kwonlyargs=[], kwonlydefaults=None, annotations={}), new=FullArgSpec(args=['input', 'dtype', 'name'], varargs=None, varkw=None, defaults=(None, None), kwonlyargs=[], kwonlydefaults=None, annotations={}), fn=<function ones_like_v2 at 0x7f3c25d2dda0>
(venv) root@c4673c623033:/#

Recreate the failure in a Python 3.11 environment using tensorflow==2.14.0 and tensorflow-probability==0.21.0, then inspect tensorflow_probability/python/internal/prefer_static.py and its _copy_docstring call for ones_like. Done means importing both tensorflow and tensorflow_probability succeeds without the Arg specs ValueError.