3.2.3　最小二乘法的梯度下降算法以及JAX实现