feature transform running issue on GPU

[ehsanngh]

Hello,

In my code, I'm using a Fourier transform like the one you used in "holography." The problem is that my code becomes very very slow on GPU while it can be executed on CPU with an acceptable runtime. I'm reaching out to see whether you faced the same issue or not.

def feature_transform(inputs):
    # Periodic BC in x
    w = 2 * np.pi / P
    x, y = w * inputs[:, :1], inputs[:, 1:]
    return tf.concat(
        (
            tf.math.cos(x),
            tf.math.sin(x),
            tf.math.cos(2 * x),
            tf.math.sin(2 * x),
            tf.math.cos(3 * x),
            tf.math.sin(3 * x),
            tf.math.cos(4 * x),
            tf.math.sin(4 * x),
            tf.math.cos(5 * x),
            tf.math.sin(5 * x),
            tf.math.cos(6 * x),
            tf.math.sin(6 * x),
            y,
            tf.math.cos(OMEGA * y),
            tf.math.sin(OMEGA * y),
        ),
        axis=1,
    )

To explain more about this issue, I found that the tensor slicing like this inputs[:, :1] and the function tf.concat are making the code slow. I just want to know whether you faced this issue before. Thanks for your time.

[ehsanngh]

Update: Switching the backend from tf2 to tf1 will solve the issue.

[lululxvi]

https://github.com/lululxvi/deepxde/discussions/1006