import numpy as np

class Dropout:
    def __init__(self, dropout_rate=0.5):
        self.drop_rate = dropout_rate
        self.mask = None

    def forward(self, x, is_train=True):
        if is_train:
            # x必须是numpy数组才有shape
            random_mat = np.random.rand(*x.shape)
            self.mask = random_mat > self.drop_rate
            output = x * self.mask
            return output
        else:
            scale = 1 - self.drop_rate
            output = x * scale
            return output

    def backward(self, upstream_grad):
        grad_x = upstream_grad * self.mask
        return grad_x

# 测试代码
if __name__ == "__main__":
    dropout_layer = Dropout(dropout_rate=0.5)
    # 正确：传入numpy数组，不能传纯数字
    x_data = np.array([10, 20, 30, 40])
    # 训练模式前向传播
    train_out = dropout_layer.forward(x_data, is_train=True)
    print("训练输出：", train_out)
    # 推理模式前向传播
    pred_out = dropout_layer.forward(x_data, is_train=False)
    print("推理输出：", pred_out)
    # 反向梯度测试
    grad = dropout_layer.backward(np.array([1,1,1,1]))
    print("反向梯度：", grad)