JAX is a fast backend simulator, but you don't necessarily have to use JAX. I think default.qubit and lightening.qubit are both installed when you install pennylane.

So if you set

dev = qml.device("default.qubit", wires=N)

I think it won't require JAX.

btw, I didn't use JAX in the code either.