Matplotlib绘制最小二乘MSE等高线图及形状不兼容报错解决

报错本质是矩阵点乘的维度不兼容：传入的w1是50*50的二维网格矩阵，和一维数据集x执行dot操作时，numpy/jax会严格按照矩阵乘法规则校验维度，最终触发形状不匹配的错误。
线性回归的预测逻辑本质是元素级运算，不需要依赖矩阵点乘，强行使用dot操作，在传入高维参数网格时就会必然出现维度冲突。

利用numpy/jax内置的广播机制实现全向量化计算，不需要编写任何Python层的循环，计算效率远高于逐点遍历的实现方式：

• 调整数据集维度结构：将原本一维的样本特征x、目标值t调整为(样本量, 1)的二维形状，为广播运算做准备
• 改写MSE计算逻辑：将原有的点乘操作替换为元素级乘法，聚合残差平方和时指定沿着样本维度计算，自动适配参数网格的二维形状
• 直接传入meshgrid生成的w、b参数网格，一次性计算整个参数平面上的MSE损失值

import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt

def make_mse(x, t):  
  def mse(w, b):
    # 基于广播自动匹配维度，支持传入单组参数或参数网格
    predictions = x * w + b
    # 沿着样本维度聚合损失，输出和参数网格同形状的损失值矩阵
    return jnp.sum(jnp.power(predictions - t, 2), axis=0) / 2
  return mse 

# 生成训练数据集
x = np.linspace(-1.0, 1.0, 500)
t = 5 * x + 1
# 调整维度为(样本量, 1) 适配广播规则
x = x.reshape(-1, 1)
t = t.reshape(-1, 1)

# 生成参数平面网格
xcoord = np.linspace(-10.0, 10.0, 100)
ycoord = np.linspace(-10.0, 10.0, 100)
w_grid, b_grid = np.meshgrid(xcoord, ycoord)

# 批量计算网格上所有点的MSE值
Z = make_mse(x, t)(w_grid, b_grid)

# 绘制等高线图
plt.figure(figsize=(8, 6))
contour = plt.contour(w_grid, b_grid, Z, levels=30, cmap='viridis')
plt.clabel(contour, inline=True, fontsize=10)
plt.xlabel('w (权重参数)')
plt.ylabel('b (偏置参数)')
plt.title('线性回归MSE损失等高线图')
plt.show()

• 计算效率高：全程使用底层向量化运算，比逐点遍历网格的实现速度快1~2个数量级
• 复用性强：同一套MSE函数既可以在训练时接收单组参数计算损失，也可以在绘图时接收整个参数网格批量计算，不需要修改函数逻辑
• 逻辑自洽：完全匹配线性回归的数学定义，绘制的等高线会自然呈现凸函数的椭圆形状，最低点精准对应数据集的真实参数w=5, b=1，可以直接叠加梯度下降的迭代路径完成可视化

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