Very different numerical results on GPU vs CPU using mmd

[alec-hoyland]

When using mmd with CUDA, I am getting wildly different numerical results than on CPU. Is this a bug or am I doing something wrong?

julia> using Flux, IPMeasures

julia> A = f32(rand(200, 128)); B = f32(rand(200, 128));

julia> A2 = gpu(A); B2 = gpu(B);

julia> mmd(GaussianKernel(2), A, B)
-2.9192196f-7

julia> mmd(GaussianKernel(2), A2, B2)
-2.5343708f-8

julia> mmd(GaussianKernel(80), A2, B2)
-1.0104629f-6

julia> mmd(GaussianKernel(80), A, B)
-1.1662798f-5

julia> mmd(GaussianKernel(80), 100*A2, B2)
29872.396f0

julia> mmd(GaussianKernel(80), 100*A, B)
-0.0033887483f0

julia>

My system

Julia Version 1.9.4
Commit 8e5136fa297 (2023-11-14 08:46 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 24 × AMD Ryzen 9 7900X 12-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, znver3)
  Threads: 1 on 24 virtual cores

CUDA toolkit 11.7, artifact installation
NVIDIA driver 545.29.6, for CUDA 12.3
CUDA driver 12.3