# Fast diag(A'\*B)
**URL:** https://discourse.julialang.org/t/fast-diag-a-b/29363
**Category:** Numerics
**Tags:** question
**Created:** [October 1, 2019, 12:40pm UTC](https://discourse.julialang.org/t/fast-diag-a-b/29363 "2019-10-01T12:40:03Z")
**Posts on this page:** 7
**Page:** 1
### Author:  [@Tamas\_Papp](https://discourse.julialang.org/u/Tamas_Papp)
#### Post date: [October 1, 2019, 12:40pm UTC](https://discourse.julialang.org/t/fast-diag-a-b/29363/1 "2019-10-01T12:40:03Z")
Is there a fast way to compute dot products of the columns of two similar matrices? MWE:
```julia
using LinearAlgebra, BenchmarkTools
f1(A, B) = diag(A'*B) # suboptimal, but surprisingly fast
f2(A, B) = map(dot, eachcol(A), eachcol(B)) # map
A = randn(150, 25)
B = randn(size(A)...)
@belapsed f1(A, B) # 9.2 μs
@belapsed f2(A, B) # 1.5 μs
```
I can write a loop, but I was kind of hoping there is something in BLAS already, but could not find anything.
---
### Author:  [@improbable22](https://discourse.julialang.org/u/improbable22)
#### Post date: [October 1, 2019, 2:11pm UTC](https://discourse.julialang.org/t/fast-diag-a-b/29363/2 "2019-10-01T14:11:35Z")
This isn’t quite what you want, but since it’s faster than `f1` they must be doing something right…
```julia
using Distances
@belapsed colwise(CosineDist(), A, B) # 3.2 μs
```
---
### Author:  [@ericphanson](https://discourse.julialang.org/u/ericphanson)
#### Post date: [October 1, 2019, 6:17pm UTC](https://discourse.julialang.org/t/fast-diag-a-b/29363/3 "2019-10-01T18:17:16Z")
You can do `sum(A .* B; dims = 1)` assuming the `’` is a transpose (otherwise you need an extra complex conjugate on `A`). This is just that (A’B)\_{ii} = \sum\_j (A’)\_{ij} B\_{ji} = \sum\_j A\_{ji} B\_{ji} in the real case.
Edit: I’m on my phone, otherwise I’d check it numerically and benchmark. But maybe it’s fast?
Edit2: fixed index in sums
---
### Author:  [@Tamas\_Papp](https://discourse.julialang.org/u/Tamas_Papp)
#### Post date: [October 1, 2019, 6:30pm UTC](https://discourse.julialang.org/t/fast-diag-a-b/29363/4 "2019-10-01T18:30:39Z")
That forms an intermediate matrix, and is much worse than `map(dot, ...)`:
```julia
f4(A, B) = vec(sum(A .* B; dims = 1))
@belapsed f4(A, B) # 3.5 μs
```
If there is no BLAS for this, then I guess I will stick with `map(dot, ...)`.
---
### Author:  [@Per](https://discourse.julialang.org/u/Per)
#### Post date: [October 1, 2019, 6:32pm UTC](https://discourse.julialang.org/t/fast-diag-a-b/29363/5 "2019-10-01T18:32:49Z")
I don’t think there’s much optimization that can be done here, other that what’s already done in `dot`. You could make the loop over columns multi-threaded, but this is going to be memory bound anyways, so…
---
### Author:  [@ericphanson](https://discourse.julialang.org/u/ericphanson)
#### Post date: [October 1, 2019, 6:33pm UTC](https://discourse.julialang.org/t/fast-diag-a-b/29363/6 "2019-10-01T18:33:36Z")
(Thanks fixing the index in your reply). If you need to do this many times you could preallocate that intermediary. Not sure if that would beat the map option though. I’m not a BLAS expert by any means, but that trick has been useful in the full-trace case for me (full sum instead of row sums).
---
### Author:  [@Mason](https://discourse.julialang.org/u/Mason)
#### Post date: [October 1, 2019, 9:05pm UTC](https://discourse.julialang.org/t/fast-diag-a-b/29363/7 "2019-10-01T21:05:29Z")
I can’t find a solution that outperforms `map(dot, ...)` here on runtime, but if you’re worried about allocations, one option would be this:
```julia
using TensorCast
f3(A, B) = @reduce C[i] := sum(j) A'[i, j] * B[j, i] lazy
A = randn(150, 25)
B = randn(size(A)...)
@btime f1($A, $B) # 9.777 μs (2 allocations: 5.34 KiB)
@btime f2($A, $B) # 1.475 μs (58 allocations: 2.78 KiB)
@btime f3($A, $B) # 8.523 μs (16 allocations: 688 bytes)
```
so you end up trading ~6x runtime speed for ~1/4 the memory usage. Often not worth it. The runtime overhead of TensorCast here goes down to ~`4µs` if you delete the `lazy` option, but the memory usage also goes up quite a bit.
Depending on your actual use-case, if you’re spending a lot of time in the GC then maybe TensorCast.jl is useful here. The other thing to consider is if preallocating will help you here.