Just a TODO that SimpleShuffleLayer and BroadcastJoinLayer both use dumps_task in their __dask_distributed_unpack__ methods, which leads to pickling on the scheduler, which we want to avoid.

Additionally, could this manual use of pickle.dumps in BroadcastJoinLayer.__dask_distributed_pack__ lead to a double-pickle scenario, where when the task actually gets executed, the kwarg values are bytestrings, not the actual unpickled values:

Related to dask/distributed#4699.

cc @rjzamora @madsbk