Repository: dask/dask-ml

Issue:
add squared arg to mean_squared_error
Was going though @jacobtomlinson dask tutorial and I saw

```
from dask_ml.metrics import mean_squared_error
from math import sqrt

sqrt(mean_squared_error(y_test, y_predicted))
```

Wonder if it's possible to add the squared arg - same as scikit-learn (https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html; https://github.com/scikit-learn/scikit-learn/blob/fd237278e/sklearn/metrics/_regression.py#L266)

May be a one line addition?

https://github.com/dask/dask-ml/blob/e5a2f2735be3f52b85dcb5d81d43ce0e5449becb/dask_ml/metrics/regression.py#L32


Discussion:
TomAugspurger: Thanks, just one small comment.

The CI failures can be ignored.

Patch:
diff --git a/dask_ml/metrics/regression.py b/dask_ml/metrics/regression.py
index afb98aa1..a72bb353 100644
--- a/dask_ml/metrics/regression.py
+++ b/dask_ml/metrics/regression.py
@@ -34,6 +34,7 @@ def mean_squared_error(
     y_pred: ArrayLike,
     sample_weight: Optional[ArrayLike] = None,
     multioutput: Optional[str] = "uniform_average",
+    squared: bool = True,
     compute: bool = True,
 ) -> ArrayLike:
     _check_sample_weight(sample_weight)
@@ -48,6 +49,8 @@ def mean_squared_error(
     else:
         raise ValueError("Weighted 'multioutput' not supported.")
     result = output_errors.mean()
+    if not squared:
+        result = da.sqrt(result)
     if compute:
         result = result.compute()
     return result
diff --git a/docs/source/contributing.rst b/docs/source/contributing.rst
index 99cbb8f5..c4df2d1f 100644
--- a/docs/source/contributing.rst
+++ b/docs/source/contributing.rst
@@ -29,7 +29,7 @@ in the ``ci`` directory.
 
 .. code-block:: none
 
-   conda env create -f ci/environment-3.7.yaml --name=dask-ml-dev
+   conda env create -f ci/environment-latest.yaml --name=dask-ml-dev
 
 to create a conda environment and install all the dependencies.
 
@@ -43,9 +43,14 @@ Building dask-ml
 Dask-ML is a pure-python repository. Development installation should be as simple as
 cloning the repository and running the following in the cloned directory:
 
-.. code-block:: none
+``pip``::
+
+  python -m pip install -e ".[dev]"
 
-   python -m pip install -e ".[dev]"
+``conda``::
+   
+  conda activate dask-ml-dev
+  python -m pip install --no-deps -e .
 
 If you have any trouble, please open an issue on the
 `dask-ml issue tracker <https://github.com/dask/dask-ml/issues>`_.


Test patch:
diff --git a/ci/environment-latest.yaml b/ci/environment-latest.yaml
new file mode 120000
index 00000000..f191f663
--- /dev/null
+++ b/ci/environment-latest.yaml
@@ -0,0 +1,1 @@
+environment-3.8.yaml
\ No newline at end of file
diff --git a/tests/metrics/test_regression.py b/tests/metrics/test_regression.py
index 42162bb0..c2cbeb0b 100644
--- a/tests/metrics/test_regression.py
+++ b/tests/metrics/test_regression.py
@@ -35,3 +35,16 @@ def test_ok(metric_pairs, compute):
         assert isinstance(result, da.Array)
     expected = m2(a, b)
     assert abs(result - expected) < 1e-5
+
+
+@pytest.mark.parametrize("squared", [True, False])
+def test_mse_squared(squared):
+    m1 = dask_ml.metrics.mean_squared_error
+    m2 = sklearn.metrics.mean_squared_error
+
+    a = da.random.uniform(size=(100,), chunks=(25,))
+    b = da.random.uniform(size=(100,), chunks=(25,))
+
+    result = m1(a, b, squared=squared)
+    expected = m2(a, b, squared=squared)
+    assert abs(result - expected) < 1e-5