Dask-Cudf Matrix Multiplication ValueError: Axis dimension mismatch

I tried to perform matrix multiplication with dask that throws an Axis dimension mismatch.
Also I tried same replication with CUDF that works fine.

please find the notebook attached for better stimulation.
daskcudf_simulated_jono.ipynb (63.4 KB)

daskcudf_simulated_jono.ipynb (63.4 KB)

For more context here as @yash.kushwaha also part of my team, I have also tried this on simulated data (please check the attached Jupyter notebook). The aim is to perform matrix multiplication on dask-cuDF to utilise GPU for our recommendation operations. As we are exploring this, we did encounter out of memory issue (currently exploring the set-up of dask kubernetes) on GPU utilisation.

Aside of that issue, now, we took a small sample (simulated) of what we have in our real environment case. We are able to perform cpu based matrix multiplication (current baseline), dask matrix multiplication, and normal cuDF matrix multiplication.

However, if you check the cell #19, it fails with the error: “axis dimension mismatch” although we have included the required meta.