I’ve tested multiple packages to compute the gradient of a log-likelihood (huge sum of the same function applied to different data points): const AV = AbstractVector{T} where T normal_pdf(x::Real, mean::Real, var::Real) = exp(-(x - mean)^2 / (2var)) / sqrt(2π * var) function mixture_loglikeli…

Any faster way of computing small gradients?

baggepinnen August 9, 2022, 3:04pm 4

Isn’t Jax using multithreading by default? If so, compare with that in Julia as well. What data type is Jax using? Could they be defaulting to Float32?

See, e.g thus thread for similar benchmarks