让 TensorFlow 对自定义类使用双精度

问：

请考虑以下 TensorFlow 代码，该代码不使用 Keras。我希望 TensorFlow 以双精度执行计算。对于该函数，我能够弄清楚如何将参数传递给装饰器。现在，如果我将变量放在一个类中并定义我的乘法函数，我该如何应用，以便以双精度进行计算？tfmulttf.functionDatadatamulttf.functiondatamult

import tensorflow as tf

# this doesn't help - because I'm not using keras
# tf.keras.backend.set_floatx('float64')

aa = 1.00006
bb = 1.00007

def simplemult(aa,bb):
    return aa*bb

# I can use arguments to the decorator to get tf to treat the arguments as float64
@tf.function(input_signature=[tf.TensorSpec(shape=None,dtype=tf.float64),tf.TensorSpec(shape=None,dtype=tf.float64)])
def tfmult(aa,bb):
    return aa*bb

# Now let us put the numbers inside a class
class Data():
    def __init__(self,xx):
        self.xx = xx
       
p1 = Data(aa)
p2 = Data(bb)

# How do I apply a tf.function decorator to this function so that operations are done in double precision
# @tf.function - defaults to single precision
@tf.function
def datamult(p1,p2):
    return (p1.xx)*(p2.xx)

print(f'{tfmult(aa,bb)=}')
print(f'{simplemult(aa,bb)=}')
print(f'{datamult(p1,p2)=}')