How to solve the backpropagation problem of tf.signal.rfft?
Hi @Evan-0715, the code provided above does not raise any error. Please find the gist [here](https://colab.sandbox.google.com/gist/synandi/03334bfe2386f6dc345e3e063a7d7d57/60163.ipynb). Kindly share complete reproducible code so that we can replicate the issue from our end. Thank you!  ```
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models, Model, Sequential
from keras.layers import concatenate
from tensorflow.keras.layers import Conv2D, Flatten, Dense
from tensorflow import matmul, reshape, reduce_sum, transpose

def hilbert_transform(x, N=None, axis=-1):
    if x.dtype == tf.complex64 or x.dtype == tf.complex128:
        raise ValueError("x must be real.")
    if N is None:
        N = x.shape[axis]
    if N <= 0:
        raise ValueError("N must be positive.")
    x_br = tf.signal.rfft3d(x[:, :, :1])
    x_bi = tf.signal.rfft3d(x[:, :, 1:])
    x_complex = concatenate([x_br, x_bi], axis=-1)

    h = np.zeros(N)
    if N % 2 == 0:
        h[0] = h[N // 2] = 1
        h[1:N // 2] = 2
    else:
        h[0] = 1
        h[1:(N + 1) // 2] = 2
    if x.shape.ndims > 1:
        ind = [tf.newaxis] * x.shape.ndims
        ind[axis] = slice(None)
        h = h[tuple(ind)]
    x_complex_h = x_complex * h
    x_bj = tf.signal.ifft3d(x_complex_h[:, :, :1])
    x_bk = tf.signal.ifft3d(x_complex_h[:, :, 1:])
    x_hilbert = concatenate([tf.math.imag(x_bj), tf.math.imag(x_bk)], axis=-1)
    return x_hilbert
class ImageUpgradingBlock(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.conv1 = tf.keras.layers.Conv2D(filters=3, kernel_size=64)
        self.conv2 = tf.keras.layers.Conv2D(filters=3, kernel_size=64)

    def call(self, inputs, *args, **kwargs):
        maps_x = inputs[:, :, :, :3]
        v1 = self.conv1(maps_x)
        v2 = self.conv2(maps_x)
        v1 = reshape(v1, [-1, 3, 1])  # basis vector with shape(-1,3,1)
        v2 = reshape(v2, [-1, 3, 1])
        V = concatenate([v1, v2], axis=-1)  # the 2-D column subspace with shape(_,3,2)
        V = V / (1e-6 + reduce_sum(abs(V), axis=1, keepdims=True))
        b_, h_, w_, c = maps_x.shape
        maps_x_t = transpose(reshape(maps_x, [-1, h_ * w_, c]), perm=[0, 2, 1])
        V_t = transpose(V, perm=[0, 2, 1])  
        mat = matmul(V_t, V)
        mat_inv = tf.linalg.inv(mat)
        # Projection --->  Y = V×(V_t×V)_(-1)×V_t×(X)
        reconstruct_x_t = matmul(V, matmul(matmul(mat_inv, V_t), maps_x_t))
        if b_ is None:  
            reconstruct_x = tf.linalg.lstsq(tf.squeeze(V, axis=[0]), tf.squeeze(reconstruct_x_t, axis=[0]))
            reconstruct_x = tf.expand_dims(reconstruct_x, axis=0)  # shape(1,c,h*w) why can‘t(None,c,h*w)
        else:
            reconstruct_x = tf.linalg.lstsq(V, reconstruct_x_t)
        reconstruct_x = transpose(reconstruct_x, perm=[0, 2, 1])  # shape(-1,h*w,c)
        reconstruct_x_h = hilbert_transform(reconstruct_x, axis=1)
        reconstruct_x = reshape(concatenate([reconstruct_x, reconstruct_x_h], axis=-1), [-1, h_, w_, c * 2 - 2])
        return reconstruct_x

    def get_config(self):
        config = super().get_config()
        return config
def RQIUNet():
    input_image = layers.Input(shape=(64, 64, 3), dtype="float32")
    x_4d = ImageUpgradingBlock()(input_image)
    result = Conv2D(3, 3, padding="same")(x_4d)
    model = models.Model(inputs=input_image, outputs=result)
    return model

loss_fn = tf.keras.losses.MeanSquaredError()
optimizer = tf.keras.optimizers.Adam()


@tf.function
def train_step(inputs, targets):
    with tf.GradientTape() as tape:
        predictions = RQIUNet()(inputs)
        loss = loss_fn(targets, predictions)
    gradients = tape.gradient(loss, RQIUNet().trainable_variables)
    optimizer.apply_gradients(zip(gradients, RQIUNet().trainable_variables))
    return loss


x = tf.random.normal((1, 64, 64, 3))
y = tf.random.normal((1, 64, 64, 3))

for i in range(10):
    loss = train_step(x, y)
    print(f"Step {i}, Loss: {loss:.4f}")
```
Hi, synandi. As the original project was quite large, I converted it into a small instance, which took a little time, sorry for that. If you run this line of code, you will get the error:
    LookupError: gradient registry has no entry for: RFFT3D
Waiting for your reply, it will help me a lot in my experiment. @Evan-0715 Thank you for sharing the complete reproducible code. I was able to replicate the error in Colab using TF v2.12 and tf-nightly(v2.13.0a20230330). Kindly refer to this [gist-2.12](https://colab.sandbox.google.com/gist/synandi/57f3d2e6c05981c76d519d44f334ef59/60163_3.ipynb) and [tf-nightly](https://colab.sandbox.google.com/gist/synandi/f9aa87c44146cdaee92ccf7e739fab0a/60163_nightly.ipynb). Thank you!

 @Evan-0715 You are right that TensorFlow might not have implemented gradients for tf.signal.rfft3d. One way to resolve this is to replace the rfft3d and ifft3d functions with their 1D counterparts, rfft and irfft, which have gradients defined.

Here's the modified hilbert_transform function using tf.signal.rfft and tf.signal.irfft:  

def hilbert_transform(x, N=None, axis=-1):
    if x.dtype == tf.complex64 or x.dtype == tf.complex128:
        raise ValueError("x must be real.")
    if N is None:
        N = x.shape[axis]
    if N <= 0:
        raise ValueError("N must be positive.")
    x_br = tf.signal.rfft(x[:, :, :1])
    x_bi = tf.signal.rfft(x[:, :, 1:])
    x_complex = tf.concat([x_br, x_bi], axis=-1)

    h = np.zeros(N)
    if N % 2 == 0:
        h[0] = h[N // 2] = 1
        h[1:N // 2] = 2
    else:
        h[0] = 1
        h[1:(N + 1) // 2] = 2
    if x.shape.ndims > 1:
        ind = [tf.newaxis] * x.shape.ndims
        ind[axis] = slice(None)
        h = h[tuple(ind)]
    x_complex_h = x_complex * h
    x_bj = tf.signal.irfft(x_complex_h[:, :, :1])
    x_bk = tf.signal.irfft(x_complex_h[:, :, 1:])
    x_hilbert = tf.concat([tf.math.imag(x_bj), tf.math.imag(x_bk)], axis=-1)
    return x_hilbert

However, there is another issue in your code. You are creating a new instance of the model at every training step, which will reset the model's weights and gradients. You should create an instance of the model and pass it as an argument to the train_step function.

Here's the modified train_step function:

@tf.function
def train_step(model, inputs, targets):
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss = loss_fn(targets, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

And you should create the model outside the loop:
x = tf.random.normal((1, 64, 64, 3))
y = tf.random.normal((1, 64, 64, 3))

model = RQIUNet()

for i in range(10):
    loss = train_step(model, x, y)
    print(f"Step {i}, Loss: {loss:.4f}")


With these changes, the code should run without the LookupError and should train the model correctly Thanks a lot! @synandi @Adesoji1 It really helped me to implement the Hilbert transform. If a input has tensor shape (b,h,w,2), it can be considered as a complex signal. Tensorflow's Hilbert transform code for its implementation is:
```
def hilbert_transform(x, N=None, axis=-1):
    ##Constructing complex number
    x_complex = tf.complex(x[:, :, :1], x[:, :, 1:])
    if x_complex.dtype != tf.complex64 and x_complex.dtype != tf.complex128:
        raise ValueError("x must be complex.")
    if N is None:
        N = x.shape[axis]
    if N <= 0:
        raise ValueError("N must be positive.")
    x_t = transpose(x_complex, perm=[0, 2, 1])
    x_f = tf.signal.fft(x_t)
    x_f = transpose(x_f, perm=[0, 2, 1])

    h = np.zeros(N)
    if N % 2 == 0:
        h[0] = h[N // 2] = 1
        h[1:N // 2] = 2
    else:
        h[0] = 1
        h[1:(N + 1) // 2] = 2
    if x.shape.ndims > 1:
        ind = [tf.newaxis] * x.shape.ndims
        ind[axis] = slice(None)
        h = h[tuple(ind)]
    x_f_h = x_f * h
    x_f = transpose(x_f_h, perm=[0, 2, 1])
    x_hilbert = tf.signal.ifft(x_f)
    x_hilbert = transpose(x_hilbert, perm=[0, 2, 1])
    x_hilbert = concatenate([tf.math.imag(x_hilbert), tf.math.real(x_hilbert)], axis=-1)
    return x_hilbert
```
It will return the output of a complex linear time-invariant system which has the shape (b,h,w,2). Thanks again! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60163">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60163">No</a>
 In addition, this can be added to the network for forward and backward propagation.
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Feature Request

### Have you reproduced the bug with TF nightly?

No

### Source

source

### Tensorflow Version

tf 2.11

### Custom Code

Yes

### OS Platform and Distribution

Windows 10

### Mobile device

_No response_

### Python version

3.9

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

```shell
I x wanted to implement the Hilbert transform in the model built by tensorflowd, but I found no callable API. so I implemented it myself using tf.signal.rfft and tf.signal.ifft. However, when I was training, I found the following error reported:
LookupError: gradient registry has no entry for: RFFT3D
I think it may be that tensorflow has not defined the corresponding backpropagation method for rfft3d yet. I want to provide my code train.py  &  model.py
```


### Standalone code to reproduce the issue

```
def hilbert_transform(x, N=None, axis=-1):
    if x.dtype == tf.complex64 or x.dtype == tf.complex128:
        raise ValueError("x must be real.")
    if N is None:
        N = x.shape[axis]
    if N <= 0:
        raise ValueError("N must be positive.")
    x_br = tf.signal.rfft3d(x[:, :, :1])
    x_bi = tf.signal.rfft3d(x[:, :, 1:])
    x_complex = concatenate([x_br, x_bi], axis=-1)

    h = np.zeros(N)
    if N % 2 == 0:
        h[0] = h[N // 2] = 1
        h[1:N // 2] = 2
    else:
        h[0] = 1
        h[1:(N + 1) // 2] = 2
    if x.shape.ndims > 1:
        ind = [tf.newaxis] * x.shape.ndims
        ind[axis] = slice(None)
        h = h[tuple(ind)]
    x_complex_h = x_complex * h
    x_bj = tf.signal.ifft3d(x_complex_h[:, :, :1])
    x_bk = tf.signal.ifft3d(x_complex_h[:, :, 1:])
    x_hilbert = concatenate([tf.math.imag(x_bj), tf.math.imag(x_bk)], axis=-1)
    return x_hilbert
```


### Relevant log output

_No response_</details>