Probably not your issue, but one kind of annoying bit is that the inputs need to be tensors. I ended up calling partial on the function I was messing around with and then vmapping the partial, which seemed to work.