jax-ldt — pure-JAX (Hyperplane) Decision Trees

Overview

jax-ldt is a JAX rewrite of LLNL's PyTorch hyperplanetree. It builds piecewise-linear regression surrogates: decision trees where every leaf holds a linear model, and (in the hyperplane variant) splits can follow oblique directions, not just axis-aligned cuts.

Key differences from upstream: