Logistic Regression Fails with ValueError: Shapes Not Aligned

Inizia da dask_ml/linear_model/glm.py e segui il traceback fino a dask_glm/algorithms.py, families.py e utils.py, concentrandoti sulle chiamate a local_update e dot. Riproduci il fallimento con cinque o più partizioni utilizzando un dataset pubblico oppure isola i dettagli mancanti per la riproduzione; il lavoro è completato quando sono state identificate una causa confermata e un test di regressione o una risoluzione documentata.

Describe the issue:
I was originally getting the exact same error message as Jreyno40 on #770. I then upgraded dask and dask-ml and ended up getting a different error when fitting the logistic regression.

Background on original AttributeError:
I am trying to train a logistic regression on a dask array using dask-ml. I am testing my code by pulling my data in batches of 1000 rows. When I pull 2-4 batches of 1000 rows and load them into separate partitions of a dask dataframe (using from_delayed) I am able to fit a model successfully (spread across 5 workers). When I expand this to 5 batches of 1000 rows (still spread across 5 workers) I receive an error when the model is fitting saying: "AttributeError: 'tuple' object has no attribute 'shape'".

It's only when I start to distribute the data to more workers that I receive this error not when my partitions increase in size which is very strange (for example I could do two partitions of 50,000 rows each with no problem but as soon as there are 5 partitions or more I got the same error). I also tried increasing the number of workers to 6 and and still using 5 batches of 1000 rows as well as 6 batches of 1000 rows. Both failed with this same error. My stack trace is identical to Jreyno40's.

Environment:
dask version: 2021.11.2
dask-ml version: 1.8.0

Shapes Not Aligned Error:
I tried upgrading the versions of dask and dask-ml (see below) on the client and workers and re-running my code. I no longer receive the attribute error mentioned above but now I receive a different error on the fitting of the logistic regression: "ValueError: shapes (0,) and (739,) not aligned: 0 (dim 0) != 739 (dim 0)". I have printed out the shapes of both arrays before fitting and their shapes do align and neither of them have zero dimensions.

It seems like there might be some sort of communication issue with the dask workers or something else that I do not understand. This code also works perfectly fine with less than 5 partitions of data. Even if I expand to 10 workers my code will have the same issues where it works fine with 2-4 partitions of data and fails once there is 5.

Environment:
dask version: 2022.2.0
dask-ml version: 2022.5.27
python: version: 3.7.13
numpy version: 1.19.5

Stack Trace

  File "xy_querying_data.py", line 346, in run
    lr.fit(X_train_arr, y_train_arr)
  File "/opt/conda/lib/python3.7/site-packages/dask_ml/linear_model/glm.py", line 188, in fit
    self._coef = algorithms._solvers[self.solver](X, y, **solver_kwargs)
  File "/opt/conda/lib/python3.7/site-packages/dask_glm/utils.py", line 26, in normalize_inputs
    out = algo(Xn, y, *args, **kwargs).copy()
  File "/opt/conda/lib/python3.7/site-packages/dask_glm/algorithms.py", line 265, in admm
    new_betas = np.array(da.compute(*new_betas))
  File "/opt/conda/lib/python3.7/site-packages/dask/base.py", line 573, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/distributed/client.py", line 2994, in get
    results = self.gather(packed, asynchronous=asynchronous, direct=direct)
  File "/opt/conda/lib/python3.7/site-packages/distributed/client.py", line 2152, in gather
    asynchronous=asynchronous,
  File "/opt/conda/lib/python3.7/site-packages/distributed/utils.py", line 310, in sync
    self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
  File "/opt/conda/lib/python3.7/site-packages/distributed/utils.py", line 376, in sync
    raise exc.with_traceback(tb)
  File "/opt/conda/lib/python3.7/site-packages/distributed/utils.py", line 349, in f
    result = yield future
  File "/opt/conda/lib/python3.7/site-packages/tornado/gen.py", line 769, in run
    value = future.result()
  File "/opt/conda/lib/python3.7/site-packages/distributed/client.py", line 2009, in _gather
    raise exception.with_traceback(traceback)
  File "/opt/conda/lib/python3.7/site-packages/dask/utils.py", line 39, in apply
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/dask_glm/algorithms.py", line 300, in local_update
    maxfun=250)
  File "/opt/conda/lib/python3.7/site-packages/scipy/optimize/lbfgsb.py", line 198, in fmin_l_bfgs_b
    **opts)
  File "/opt/conda/lib/python3.7/site-packages/scipy/optimize/lbfgsb.py", line 308, in _minimize_lbfgsb
    finite_diff_rel_step=finite_diff_rel_step)
  File "/opt/conda/lib/python3.7/site-packages/scipy/optimize/optimize.py", line 262, in _prepare_scalar_function
    finite_diff_rel_step, bounds, epsilon=epsilon)
  File "/opt/conda/lib/python3.7/site-packages/scipy/optimize/_differentiable_functions.py", line 140, in __init__
    self._update_fun()
  File "/opt/conda/lib/python3.7/site-packages/scipy/optimize/_differentiable_functions.py", line 233, in _update_fun
    self._update_fun_impl()
  File "/opt/conda/lib/python3.7/site-packages/scipy/optimize/_differentiable_functions.py", line 137, in update_fun
    self.f = fun_wrapped(self.x)
  File "/opt/conda/lib/python3.7/site-packages/scipy/optimize/_differentiable_functions.py", line 134, in fun_wrapped
    return fun(np.copy(x), *args)
  File "/opt/conda/lib/python3.7/site-packages/dask_glm/algorithms.py", line 234, in wrapped
    return func(beta, X, y) + (rho / 2) * np.dot(beta - z + u,
  File "/opt/conda/lib/python3.7/site-packages/dask_glm/families.py", line 31, in pointwise_loss
    return Logistic.loglike(Xbeta, y)
  File "/opt/conda/lib/python3.7/site-packages/dask_glm/families.py", line 24, in loglike
    return (Xbeta + log1p(enXbeta)).sum() - dot(y, Xbeta)
  File "/opt/conda/lib/python3.7/site-packages/multipledispatch/dispatcher.py", line 278, in __call__
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/dask_glm/utils.py", line 126, in dot
    return np.dot(A, B)
  File "<__array_function__ internals>", line 6, in dot
ValueError: shapes (0,) and (739,) not aligned: 0 (dim 0) != 739 (dim 0)

Minimal Complete Verifiable Example:
Unfortunately I cannot create a reproducible example. The code is confidential and I am creating the dask cluster using Iguazio/MLRun which is something others wouldn't have access to. I tried creating an example with the iris dataset similar to what Jreyno40 described but everything executed fine using a local cluster (trying to make it accessible for everyone).

I'm sorry that I cannot provide I reproducible example I am hoping that some of the context that I have provided would help to understand the issue better. Please let me know if there is any other information I can provide on this error. Thank you in advance for any help you can offer!