jax

3c5ce13b - add more reasonable linearize sketch