适配TensorFlow张量的傅里叶级数函数改写需求

需要开发支持TensorFlow张量作为y值的傅里叶级数函数，当前版本仅适配浮点数，需改写以支持如下结构的输入：

BATCH_SIZE = 16
MAX_LENGTH = 50
D_MODEL = 32

input_tensor_1 = tf.random.normal((BATCH_SIZE, MAX_LENGTH, D_MODEL))
input_tensor_2 = tf.random.normal((BATCH_SIZE, MAX_LENGTH, D_MODEL))

y = [input_tensor_1, input_tensor_2]
x = tf.range(len(y), dtype=tf.float32)

原始浮点数版本代码 #

import tensorflow as tf
import numpy as np

def construct_periodic_function(x, y, num_coefficients):
    sorted_indices = tf.argsort(x)
    x_sorted = tf.gather(x, sorted_indices)
    y_sorted = tf.gather(y, sorted_indices)

    n = tf.shape(x_sorted)[0]
    T = x_sorted[-1] - x_sorted[0]  # Total period
    omega = 2 * np.pi / T  # Angular frequency

    a0 = tf.reduce_mean(y_sorted)
    an = []
    bn = []
    for i in range(1, num_coefficients + 1):
        an_i = 2 * tf.reduce_mean(y_sorted * tf.cos(i * omega * x_sorted))  # Cosine coefficients
        bn_i = 2 * tf.reduce_mean(y_sorted * tf.sin(i * omega * x_sorted))  # Sine coefficients
        an.append(an_i)
        bn.append(bn_i)

    an = tf.stack(an)
    bn = tf.stack(bn)

    def f(x_new):
        y_new = tf.zeros_like(x_new, dtype=tf.float32)
        y_new += a0
        for i in range(1, num_coefficients + 1):
            y_new += an[i-1] * tf.cos(i * omega * x_new)
            y_new += bn[i-1] * tf.sin(i * omega * x_new)
        return y_new

    return f

x = tf.linspace(0.0, 2 * np.pi, 50)
y = tf.sin(x)

num_coefficients = 10

f = construct_periodic_function(x, y, num_coefficients)

x_range = tf.linspace(0.0, 2 * np.pi, 1000)
y_range = f(x_range)

改写后的TensorFlow张量版本代码 #

import tensorflow as tf
import numpy as np

def construct_periodic_function(x, y, num_coefficients):
    # 将y列表堆叠为张量，形状变为(N, BATCH_SIZE, MAX_LENGTH, D_MODEL)，N为len(y)
    y_tensor = tf.stack(y, axis=0)
    # 按x排序
    sorted_indices = tf.argsort(x)
    x_sorted = tf.gather(x, sorted_indices)
    y_sorted = tf.gather(y_tensor, sorted_indices, axis=0)

    n = tf.shape(x_sorted)[0]
    T = x_sorted[-1] - x_sorted[0]  # 周期长度
    omega = 2 * np.pi / T  # 角频率

    # 计算a0：仅对x维度（第0维）取均值，保留张量的其他维度
    a0 = tf.reduce_mean(y_sorted, axis=0)
    an = []
    bn = []
    for i in range(1, num_coefficients + 1):
        # 计算三角函数项：x_sorted为(N,)，自动广播到与y_sorted匹配的形状
        cos_term = tf.cos(i * omega * x_sorted)
        sin_term = tf.sin(i * omega * x_sorted)
        # 对x维度取均值，保留BATCH、MAX_LENGTH、D_MODEL维度
        an_i = 2 * tf.reduce_mean(y_sorted * cos_term[:, tf.newaxis, tf.newaxis, tf.newaxis], axis=0)
        bn_i = 2 * tf.reduce_mean(y_sorted * sin_term[:, tf.newaxis, tf.newaxis, tf.newaxis], axis=0)
        an.append(an_i)
        bn.append(bn_i)

    # 堆叠系数：形状变为(num_coefficients, BATCH_SIZE, MAX_LENGTH, D_MODEL)
    an = tf.stack(an, axis=0)
    bn = tf.stack(bn, axis=0)

    def f(x_new):
        # x_new可能是任意形状的张量，需要将系数广播到匹配的形状
        # 扩展维度以兼容系数的形状：x_new -> (*, 1, 1, 1)
        x_expanded = x_new[..., tf.newaxis, tf.newaxis, tf.newaxis]
        # 初始化输出张量，形状与x_new扩展后加上BATCH、MAX_LENGTH、D_MODEL一致
        y_new = tf.broadcast_to(a0, tf.concat([tf.shape(x_new), tf.shape(a0)], axis=0))
        for i in range(1, num_coefficients + 1):
            cos_term = tf.cos(i * omega * x_expanded)
            sin_term = tf.sin(i * omega * x_expanded)
            # 系数与三角函数项相乘并叠加
            y_new += an[i-1] * cos_term
            y_new += bn[i-1] * sin_term
        return y_new

    return f

# 测试代码
BATCH_SIZE = 16
MAX_LENGTH = 50
D_MODEL = 32

input_tensor_1 = tf.random.normal((BATCH_SIZE, MAX_LENGTH, D_MODEL))
input_tensor_2 = tf.random.normal((BATCH_SIZE, MAX_LENGTH, D_MODEL))

y = [input_tensor_1, input_tensor_2]
x = tf.range(len(y), dtype=tf.float32)

num_coefficients = 10
f = construct_periodic_function(x, y, num_coefficients)

# 测试预测：x_new可以是单个值或张量，比如预测5个点
x_new = tf.linspace(0.0, 2.0, 5)
y_new = f(x_new)
# 输出形状应为(5, 16, 50, 32)
print(y_new.shape)

关键改动说明 #

• y的处理：将输入的y列表堆叠为张量，保留其原始的BATCH、MAX_LENGTH、D_MODEL维度，仅在第0维增加样本数N。
• 均值计算：reduce_mean明确指定对x维度（第0维）取均值，避免丢失张量的其他维度信息。
• 广播兼容：对x相关的三角函数项添加新维度，确保能与高维的y张量进行元素级乘法；预测时也对x_new扩展维度，实现系数与输入的广播匹配。
• 输出形状：保证预测输出的形状与x_new的形状和原始y的维度兼容，比如输入x_new为(K,)时，输出为(K, BATCH_SIZE, MAX_LENGTH, D_MODEL)。

内容的提问来源于stack exchange，提问作者Tyron