1. 设计一个接受输入并计算张量降维的层，它返回 $y_k = \sum_{i, j} W_{ijk} x_i x_j$。

import torch
import torch.nn as nn

class TensorReductionLayer(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(TensorReductionLayer, self).__init__()
        # 初始化权重张量 W，其形状为 (output_dim, input_dim, input_dim)
        self.W = nn.Parameter(torch.randn(output_dim, input_dim, input_dim))

    def forward(self, x):
        # x 的形状为 (batch_size, input_dim)
        batch_size, input_dim = x.shape

        # 计算 y_k = sum_{i,j} W_{ijk} x_i x_j
        # 首先计算 x_i x_j 的所有组合
        x_outer = torch.einsum('bi,bj->bij', x, x)  # 形状为 (batch_size, input_dim, input_dim)

        # 将 x_outer 和 W 进行元素级相乘并在最后一个维度上求和
        y = torch.einsum('bii,jik->bj', x_outer, self.W)  # 形状为 (batch_size, output_dim)

        return y

# 示例使用
input_dim = 4
output_dim = 3
batch_size = 2

# 创建层
layer = TensorReductionLayer(input_dim, output_dim)

# 创建输入张量
x = torch.randn(batch_size, input_dim)

# 前向传播
y = layer(x)
print(y)

tensor([[-0.2735,  0.7346,  4.8463],
        [-2.5451,  1.4819,  8.2211]], grad_fn=<ViewBackward0>)

2. 设计一个返回输入数据的傅立叶系数前半部分的层。

import torch
import torch.nn as nn

class FourierTransformLayer(nn.Module):
    def __init__(self):
        super(FourierTransformLayer, self).__init__()

    def forward(self, x):
        # 计算傅立叶变换
        x_fft = torch.fft.fft(x, dim=-1)
        
        # 仅取前半部分的傅立叶系数
        n = x.shape[-1]
        half_n = n // 2 + 1  # 前半部分包括中间点
        x_fft_half = x_fft[..., :half_n]
        
        return x_fft_half

# 示例使用
input_dim = 8
batch_size = 2

# 创建层
layer = FourierTransformLayer()

# 创建输入张量
x = torch.randn(batch_size, input_dim)

# 前向传播
y = layer(x)
print(y)

tensor([[ 4.5006+0.0000j, -2.9140-1.8181j, -2.9015-1.1578j,  1.0263+0.7136j,
         -4.4484+0.0000j],
        [-0.4223+0.0000j, -0.5647-0.7324j, -1.1352+2.9639j,  3.7225-0.3951j,
         -2.5417+0.0000j]])