Jax Experiencing`NAME ERROR`

UPDATE: This does seem to sample with pm.sample without a NAME ERROR