JZ port

Did break a few things for CUDA backend