Train error vs. Test error in linear regression by samples analysis
It occurs because your data is noisy. Suppose that your model is
$$ y = X\beta + noise. $$
Suppose that $\beta$ is recovered exactly. Since the noise is present, our model will NOT predict accurately, it will be always some noise. In other words, think that any time we predict the values, there is some random perturbation. The amount of perturbation is of the same order each time. When we compute the error, we square these perturbations, sum them up, and then take the square root, i.e., compute the norm
In your example, the size of the training set is growing, and, therefore, the test error WITHOUT averaging is growing as well -- you sum up squares of the perturbation. The size of the test set is steady, so the norm of error is about the same there.
Solution: use appropriate averaging to both train and test set.
train_loss = (1/np.sqrt(i) * tf.tensordot(tf.transpose(h(X2[:i], weights) - Y[:i]), (h(X2[:i], weights) - Y[:i]), axes=1))[0][0]
test_loss = (1/(2 * X2.shape[0]) * tf.tensordot(tf.transpose(h(X2, weights) - testY), (h(X2, weights) - testY), axes=1))[0][0]

P.S. X2 in my formulas are as defined in your cell, accordingly. I would suggest choosing a different name for the variable.