* Add `concat_with_broadcast` helper function

Use new helper in xt.concat
Co-authored-by: Ricardo <ricardo.vieira1994@gmail.com>

* block_diag dot rewrite
Co-authored-by: Ricardo <ricardo.vieira1994@gmail.com>

---------
Co-authored-by: Ricardo <ricardo.vieira1994@gmail.com>