Update axiom to use axiom-types

Signed-off-by: Dan Kubb <dan.kubb@gmail.com>