jax

89dfac4b - remove redundant step