背景与问题
我用的是哈密顿蒙特卡罗（hmc）方法 tensorflow-probability 模块探索自编概率函数的最可能状态。在我试图拟合的参数中，有一个真实三维场的傅里叶模式。
为了运行hmc，每个计算块都需要实现其梯度。然而，逆实fft的实现 tf.signal.irfft3d 默认情况下没有关联的渐变方法。
问题
有办法实现函数的梯度吗 irfft3d ? 我已经有一个运行，自我实现 irfft3d 更基本的 tensorflow 自动微分似乎是工作的，但我想 Package 的实际优化和稳定的实现 tf.signal.irfft3d 使用装饰器 @tf.custom_gradient 使自动微分工作。
我猜
傅立叶变换是线性的，这个问题在理论上是微不足道的。然而，在网格上编写傅里叶变换的雅可比矩阵在数值上是不可行的（因为它的维数将是巨大的）。幸运的是， tensorflow 只要求对输入向量计算雅可比矩阵的函数。我相信这可以有效地完成感谢快速傅立叶变换算法。不幸的是，在我看来 tensorflow 需要一个函数来计算应用于“上游梯度”的雅可比矩阵的倒数，我不明白：
https://www.tensorflow.org/api_docs/python/tf/custom_gradient?version=nightly
返回元组（y，grad\u fn）的函数f（*x），其中：
x是函数的Tensor输入序列（嵌套结构）。
y是在f到x中应用tensorflow 运算的Tensor输出的（嵌套结构）。
grad\u fn是一个带有签名g（grad\u ys）的函数，它返回与（展平的）x相同大小的Tensor列表-y中Tensor相对于x中Tensor的导数。graduys是一个与（展平的）y大小相同的Tensor序列，它保持y中每个Tensor的初始值梯度。
在纯数学意义上，向量参数向量值函数f的导数应该是它的雅可比矩阵j。在这里，我们将雅可比j表示为函数grad\u fn，它定义了当j与向量grad\u ys相乘时，j如何变换向量grad\u ys（grad\u ysj，向量雅可比积，或vjp）。矩阵的这种函数表示便于用于链式规则计算（例如，在反向传播算法中）。
按照文档中给出的尺寸和格式，我无法想象任何其他解决方案：

# !/usr/bin/env python3

# set up

import tensorflow as tf

n = 64

noise = tf.random.normal((n, n, n))
modes = tf.signal.rfft3d(noise)

# the function

@tf.custom_gradient
def irfft3d(x):
    def grad_fn(dy):
        return (tf.signal.rfft3d(dy))

    return (tf.signal.irfft3d(x), grad_fn)

# test

with tf.GradientTape() as gt:
    gt.watch(modes)
    rec_noise = irfft3d(modes)

dn_dm = gt.gradient(rec_noise, modes)

print(dn_dm)

它运行并返回：

tf.Tensor(
[[[262144.+0.j      0.+0.j      0.+0.j ...      0.+0.j      0.+0.j
        0.+0.j]
  [     0.+0.j      0.+0.j      0.+0.j ...      0.+0.j      0.+0.j
        0.+0.j]
  [     0.+0.j      0.+0.j      0.+0.j ...      0.+0.j      0.+0.j
        0.+0.j]
  ...
  [     0.+0.j      0.+0.j      0.+0.j ...      0.+0.j      0.+0.j
        0.+0.j]
  [     0.+0.j      0.+0.j      0.+0.j ...      0.+0.j      0.+0.j
        0.+0.j]
  [     0.+0.j      0.+0.j      0.+0.j ...      0.+0.j      0.+0.j
        0.+0.j]]], shape=(64, 64, 33), dtype=complex64)

我真的不能把我的心思都放在它上面。首先，这将是一个如此简单的解决方案，以至于我不理解为什么它没有在本地实现。但更重要的是，我只是迷失在 tensorflow 期望从这个自编梯度函数，我不能表达它的结果在数学上的方式，使我有意义。
有人知道路吗 tensorflow 你能帮我吗？