Why not doing 
```
return np.average(((y_pred - y_true) ** 2).mean(axis=1), weights=sample_weight)
```
?