2.2.3　基于JAX的线性回归模型的编写