Pure-JAX linear and hyperplane decision trees for piecewise-linear surrogate modeling - View it on GitHub