Tensorflow：将概率API应用于非恒定张量

Question

在我的损失函数中，变量（例如C）是非线性方程的解。我试图使用概率API tfp.math.secant_root（文档normalize()）来计算根。但是，secant_root中的示例仅接收恒定张量。

举一个更具体的例子，我的损失函数定义如下：

def loss(y,C):
    ## Obtain C as the solution to a nonlinear equation involving y
    return y * tf.pow(C,1.5)

其中y是预测变量，因此y的维度为[?,1]。我不确定如何应用tfp.math.secant_root，因为示例仅涉及恒定张量。