I would have to check what you’re doing in more detail to find the error, but it should be possible to write it like you did.
Note that your logp is not a proper density, as you didn’t had the jacobian from the exponentiation operation. This may or not matter for your use case