# Batched Hessian-Vector Product (on the GPU)

**URL:** https://discourse.julialang.org/t/batched-hessian-vector-product-on-the-gpu/130372
**Category:** GPU
**Tags:** question, cuda, autodiff
**Created:** [July 1, 2025, 2:30pm UTC](https://discourse.julialang.org/t/batched-hessian-vector-product-on-the-gpu/130372 "2025-07-01T14:30:38Z")
**Posts on this page:** 1
**Page:** 1

<div class="post-metadata">

### Author: ![eblumen](https://avatars.discourse-cdn.com/v4/letter/e/87869e/32.png) [@eblumen](https://discourse.julialang.org/u/eblumen)
#### Post date: [July 1, 2025, 2:30pm UTC](https://discourse.julialang.org/t/batched-hessian-vector-product-on-the-gpu/130372/1 "2025-07-01T14:30:38Z")

</div>

Hello,

I am trying to implement a sort of ‘batched’ HVP. Here’s the setup: consider some collection of particle positions x\_{i,\alpha,\beta,\gamma} where i = 1,\ldots, d is a coordinate index, and \alpha,\beta,\gamma index particle identity/system replicate. Let f be some scalar function of particle position which also depends on some parameters p\_\gamma which looks like f(\vec x\_{\alpha,\beta ,\gamma}; p\_\gamma). I would like to compute the quantities (in index notation):

\texttt{grad}\_{i,\alpha,\beta,\gamma} = \frac{\partial}{\partial x\_{i,\alpha,\beta,\gamma}} f(\vec x\_{\alpha,\beta,\gamma}), \\ \texttt{hvp}\_{i,\alpha,\beta,\gamma} = \sum\_{j=1}^d v\_{j,\alpha,\beta,\gamma} \frac{\partial}{\partial x\_{i,\alpha,\beta,\gamma}} \frac{\partial}{\partial x\_{j,\alpha,\beta,\gamma}} f(\vec x\_{\alpha,\beta,\gamma}),

where \vec v\_{\alpha,\beta,\gamma} are some vectors (one for each particle/replicate).  
Here’s some example code where I successfully implemented the first:

```julia
using CUDA, Zygote
n_dims = 3; n_α = 100; n_β = 100; n_γ = 100;

function f_single(x::AbstractVector{T}, μ::AbstractVector{T}) where {T<:Number}
    return exp(-sum(abs2,x - μ))
end

function f_batch(x::AbstractArray{T,4}, μ::AbstractArray{T,2}) where {T<:Number}
    return exp.(-dropdims(sum(abs2,x .- reshape(μ, size(μ,1), 1,1, size(μ,2)),dims=1),dims=1))
end

x = CUDA.randn(n_dims, n_α, n_β, n_γ); 
μs = CUDA.randn(n_dims, n_γ)
f_batch(x, μs)::AbstractArray{Float32,3}
f_single(view(x,:,1,1,1), view(μs,:,1))::Float32

Zygote.gradient(x -> sum(f_batch(x, μs)), x)[1]# CuArray{Float32, 4, CUDA.DeviceMemory} of size (n_dims, n_α, n_β, n_γ)

```

I can’t figure out how to successfully implement the HVP, though. It seems that all built-in methods essentially will compute something like:

\sum\_{j,\alpha',\beta',\gamma'} v\_{j,\alpha',\beta',\gamma'} \frac{\partial}{\partial x\_{i,\alpha,\beta,\gamma}} \frac{\partial}{\partial x\_{j,\alpha',\beta',\gamma'}} f(\vec x\_{\alpha',\beta',\gamma'})

which is _not_ what I want. I would greatly appreciate any recommendations on how to proceed. For my more complicated use case, it’s important that I run this on the GPU and that it is efficient as I will be computing with large arrays in high dimensions, thousands of times. Additionally, for my full use case, `f_batch` is typically more efficient than broadcasting `f_single`, so it would be ideal to work with `f_batch` if possible. I am open to Zygote, Enzyme, or any other package that works and is efficient.

I appreciate any advice anyone might offer. I’m really stuck! Thank you in advance!