Change the monkeypatch thing to using with contexts?

It would be nice to be able to do things like

```
with matrix_dims(["dim1", "dim3"]):
    chol = xe.linalg.cholesky(da)
    eig = xe.linalg.eig(da)
```

instead of having to use the monkeypatch trick ([currently documented](https://einstats.python.arviz.org/en/latest/api/generated/xarray_einstats.stats.get_default_dims.html)) or needing to pass the dimensions every time.