Unable to import tf_quant_finance as tff in Colab notebook: Monte_Carlo_Euler_Scheme.ipynb

import matplotlib.pyplot as plt
import numpy as np
import time

import tensorflow as tf

import QuantLib as ql

tff for Tensorflow Finance

import tf_quant_finance as tff

from IPython.core.pylabtools import figsize
figsize(21, 14) # better graph size for Colab

ValueError Traceback (most recent call last)
in <cell line: 12>()
10
11 # tff for Tensorflow Finance
---> 12 import tf_quant_finance as tff
13
14 from IPython.core.pylabtools import figsize

16 frames
/usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/internal/prefer_static.py in _copy_docstring(original_fn, new_fn)
82 new_spec = tf_inspect.getfullargspec(new_fn)
83 if original_spec != new_spec:
---> 84 raise ValueError(
85 'Arg specs do not match: original={}, new={}, fn={}'.format(
86 original_spec, new_spec, original_fn))

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 0x7d8479b6c5e0>

Start with Monte_Carlo_Euler_Scheme.ipynb and reproduce the tf_quant_finance import in Google Colab using the imports shown. Trace the TensorFlow Probability prefer_static.py error and identify the compatibility issue causing the argument-spec mismatch. Done means the notebook imports tf_quant_finance as tff without this ValueError.