bartz 0.4.1__tar.gz → 0.5.0__tar.gz

{bartz-0.4.1 → bartz-0.5.0}/PKG-INFO

@@ -1,22 +1,18 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: bartz
-Version: 0.4.1
+Version: 0.5.0
 Summary: Super-fast BART (Bayesian Additive Regression Trees) in Python
-License: MIT
 Author: Giacomo Petrillo
-Author-email: info@giacomopetrillo.com
+Author-email: Giacomo Petrillo <info@giacomopetrillo.com>
+License-Expression: MIT
+Requires-Dist: jax>=0.4.35,<1
+Requires-Dist: jaxlib>=0.4.35,<1
+Requires-Dist: numpy>=1.25.2,<3
+Requires-Dist: scipy>=1.11.4,<2
 Requires-Python: >=3.10
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
-Requires-Dist: jax (>=0.4.35,<1)
-Requires-Dist: jaxlib (>=0.4.35,<1)
-Requires-Dist: numpy (>=1.25.2,<3)
-Requires-Dist: scipy (>=1.11.4,<2)
-Project-URL: Bug Tracker, https://github.com/Gattocrucco/bartz/issues
+Project-URL: Documentation, https://gattocrucco.github.io/bartz/docs-dev
+Project-URL: Homepage, https://github.com/Gattocrucco/bartz
+Project-URL: Issues, https://github.com/Gattocrucco/bartz/issues
 Description-Content-Type: text/markdown
 
 [![PyPI](https://img.shields.io/pypi/v/bartz)](https://pypi.org/project/bartz/)
@@ -42,6 +38,7 @@ On CPU, bartz runs at the speed of dbarts (the fastest implementation I know of)
 - [Documentation (development version)](https://gattocrucco.github.io/bartz/docs-dev)
 - [Repository](https://github.com/Gattocrucco/bartz)
 - [Code coverage](https://gattocrucco.github.io/bartz/coverage)
+- [Benchmarks](https://gattocrucco.github.io/bartz/benchmarks)
 - [List of BART packages](https://gattocrucco.github.io/bartz/docs-dev/pkglist.html)
 
 ## Citing bartz
@@ -49,4 +46,3 @@ On CPU, bartz runs at the speed of dbarts (the fastest implementation I know of)
 Article: Petrillo (2024), "Very fast Bayesian Additive Regression Trees on GPU", [arXiv:2410.23244](https://arxiv.org/abs/2410.23244).
 
 To cite the software directly, including the specific version, use [zenodo](https://doi.org/10.5281/zenodo.13931477).
-

{bartz-0.4.1 → bartz-0.5.0}/README.md

@@ -21,6 +21,7 @@ On CPU, bartz runs at the speed of dbarts (the fastest implementation I know of)
 - [Documentation (development version)](https://gattocrucco.github.io/bartz/docs-dev)
 - [Repository](https://github.com/Gattocrucco/bartz)
 - [Code coverage](https://gattocrucco.github.io/bartz/coverage)
+- [Benchmarks](https://gattocrucco.github.io/bartz/benchmarks)
 - [List of BART packages](https://gattocrucco.github.io/bartz/docs-dev/pkglist.html)
 
 ## Citing bartz

{bartz-0.4.1 → bartz-0.5.0}/pyproject.toml

@@ -23,12 +23,12 @@
 # SOFTWARE.
 
 [build-system]
-requires = ["poetry-core>=2.0.0,<3.0.0"]
-build-backend = "poetry.core.masonry.api"
+requires = ["uv_build>=0.7.3,<0.8.0"]
+build-backend = "uv_build"
 
 [project]
 name = "bartz"
-version = "0.4.1"
+version = "0.5.0"
 description = "Super-fast BART (Bayesian Additive Regression Trees) in Python"
 authors = [
     {name = "Giacomo Petrillo", email = "info@giacomopetrillo.com"},
@@ -36,7 +36,6 @@ authors = [
 license = "MIT"
 readme = "README.md"
 requires-python = ">=3.10"
-repository = "https://github.com/Gattocrucco/bartz"
 packages = [
     { include = "bartz", from = "src" },
 ]
@@ -47,29 +46,35 @@ dependencies = [
     "scipy >=1.11.4,<2",
 ]
 
-[tool.poetry.urls]
-"Bug Tracker" = "https://github.com/Gattocrucco/bartz/issues"
-
-[tool.poetry.group.dev.dependencies]
-ipython = ">=8.22.2"
-matplotlib = "^3.8.3"
-appnope = "^0.1.4"
-tomli = "^2.0.1"
-packaging = "^24.0"
-xgboost = "^2.0.3"
-pre-commit = "^3.7.0"
-scikit-learn = "^1.4.2"
-matplotlib-label-lines = "^0.7.0"
-polars = {version = "^1.0.0", extras = ["pandas", "pyarrow"]}
-
-[tool.poetry.group.test.dependencies]
-coverage = "^7.4.3"
-pytest = "^8.1.1"
-
-[tool.poetry.group.docs.dependencies]
-Sphinx = ">=7.2.6,<9"
-numpydoc = "^1.6.0"
-myst-parser = "^4.0.1"
+[project.urls]
+Homepage = "https://github.com/Gattocrucco/bartz"
+Documentation = "https://gattocrucco.github.io/bartz/docs-dev"
+Issues = "https://github.com/Gattocrucco/bartz/issues"
+
+[dependency-groups]
+only-local = [
+    "appnope>=0.1.4",
+    "ipython>=8.36.0",
+    "matplotlib>=3.10.3",
+    "matplotlib-label-lines>=0.8.1",
+    "polars[pandas,pyarrow]>=1.29.0",
+    "pre-commit>=4.2.0",
+    "ruff>=0.11.9",
+    "scikit-learn>=1.6.1",
+    "tomli>=2.2.1",
+    "virtualenv>=20.31.2",
+    "xgboost>=3.0.0",
+]
+ci = [
+    "asv>=0.6.4",
+    "coverage>=7.8.0",
+    "myst-parser>=4.0.1",
+    "numpydoc>=1.8.0",
+    "packaging>=25.0",
+    "pytest>=8.3.5",
+    "pytest-timeout>=2.4.0",
+    "sphinx>=8.1.3",
+]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
@@ -81,6 +86,8 @@ addopts = [
     "--pdbcls=IPython.terminal.debugger:TerminalPdb",
     "--durations=3",
 ]
+timeout = 32
+timeout_method = "thread"  # when jax hangs, signals do not work
 
 # I wanted to use `--import-mode=importlib`, but it breaks importing submodules,
 # in particular `from . import util`.
@@ -119,3 +126,24 @@ local = [
     '/opt/hostedtoolcache/Python/*/*/lib/python*/site-packages/bartz/',
     'C:\hostedtoolcache\windows\Python\*\*\Lib\site-packages\bartz\',
 ]
+
+[tool.ruff]
+exclude = [".asv", "*.ipynb"]
+
+[tool.ruff.format]
+quote-style = "single"
+
+[tool.ruff.lint]
+select = [
+    "B", # bugbear: grab bag of additional stuff
+    "UP", # pyupgrade: fix some outdated idioms
+    "I", # isort: sort and reformat import statements
+    "F", # flake8
+]
+ignore = [
+    "B028", # warn with stacklevel = 2
+]
+
+[tool.uv]
+python-downloads = "never"
+python-preference = "only-system"

{bartz-0.4.1 → bartz-0.5.0}/src/bartz/BART.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/BART.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -27,11 +27,8 @@ import functools
 import jax
 import jax.numpy as jnp
 
-from . import jaxext
-from . import grove
-from . import mcmcstep
-from . import mcmcloop
-from . import prepcovars
+from . import grove, jaxext, mcmcloop, mcmcstep, prepcovars
+
 
 class gbart:
     """
@@ -53,10 +50,11 @@ class gbart:
         Whether to use predictors quantiles instead of a uniform grid to bin
         predictors.
     sigest : float, optional
-        An estimate of the residual standard deviation on `y_train`, used to
-        set `lamda`. If not specified, it is estimated by linear regression.
-        If `y_train` has less than two elements, it is set to 1. If n <= p, it
-        is set to the variance of `y_train`. Ignored if `lamda` is specified.
+        An estimate of the residual standard deviation on `y_train`, used to set
+        `lamda`. If not specified, it is estimated by linear regression (with
+        intercept, and without taking into account `w`). If `y_train` has less
+        than two elements, it is set to 1. If n <= p, it is set to the standard
+        deviation of `y_train`. Ignored if `lamda` is specified.
     sigdf : int, default 3
         The degrees of freedom of the scaled inverse-chisquared prior on the
         noise variance.
@@ -82,6 +80,12 @@ class gbart:
     offset : float, optional
         The prior mean of the latent mean function. If not specified, it is set
         to the mean of `y_train`. If `y_train` is empty, it is set to 0.
+    w : array (n,), optional
+        Coefficients that rescale the error standard deviation on each
+        datapoint. Not specifying `w` is equivalent to setting it to 1 for all
+        datapoints. Note: `w` is ignored in the automatic determination of
+        `sigest`, so either the weights should be O(1), or `sigest` should be
+        specified by the user.
     ntree : int, default 200
         The number of trees used to represent the latent mean function.
     numcut : int, default 255
@@ -108,6 +112,8 @@ class gbart:
         The number of iterations (including skipped ones) between each log.
     seed : int or jax random key, default 0
         The seed for the random number generator.
+    initkw : dict
+        Additional arguments passed to `mcmcstep.init`.
 
     Attributes
     ----------
@@ -135,8 +141,6 @@ class gbart:
         The number of trees.
     maxdepth : int
         The maximum depth of the trees.
-    initkw : dict
-        Additional arguments passed to `mcmcstep.init`.
 
     Methods
     -------
@@ -158,10 +162,13 @@ class gbart:
     - A lot of functionality is missing (variable selection, discrete response).
     - There are some additional attributes, and some missing.
 
-    The linear regression used to set `sigest` adds an intercept.
     """
 
-    def __init__(self, x_train, y_train, *,
+    def __init__(
+        self,
+        x_train,
+        y_train,
+        *,
         x_test=None,
         usequants=False,
         sigest=None,
@@ -173,6 +180,7 @@ class gbart:
         maxdepth=6,
         lamda=None,
         offset=None,
+        w=None,
         ntree=200,
         numcut=255,
         ndpost=1000,
@@ -180,26 +188,41 @@ class gbart:
         keepevery=1,
         printevery=100,
         seed=0,
-        initkw={},
-        ):
-
+        initkw=None,
+    ):
         x_train, x_train_fmt = self._process_predictor_input(x_train)
-
-        y_train, y_train_fmt = self._process_response_input(y_train)
+        y_train, _ = self._process_response_input(y_train)
         self._check_same_length(x_train, y_train)
+        if w is not None:
+            w, _ = self._process_response_input(w)
+            self._check_same_length(x_train, w)
 
         offset = self._process_offset_settings(y_train, offset)
         scale = self._process_scale_settings(y_train, k)
-        lamda, sigest = self._process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda, offset)
+        lamda, sigest = self._process_noise_variance_settings(
+            x_train, y_train, sigest, sigdf, sigquant, lamda, offset
+        )
 
         splits, max_split = self._determine_splits(x_train, usequants, numcut)
         x_train = self._bin_predictors(x_train, splits)
-
-        y_train = self._transform_input(y_train, offset, scale)
-        lamda_scaled = lamda / (scale * scale)
-
-        mcmc_state = self._setup_mcmc(x_train, y_train, max_split, lamda_scaled, sigdf, power, base, maxdepth, ntree, initkw)
-        final_state, burnin_trace, main_trace = self._run_mcmc(mcmc_state, ndpost, nskip, keepevery, printevery, seed)
+        y_train, lamda_scaled = self._transform_input(y_train, lamda, offset, scale)
+
+        mcmc_state = self._setup_mcmc(
+            x_train,
+            y_train,
+            w,
+            max_split,
+            lamda_scaled,
+            sigdf,
+            power,
+            base,
+            maxdepth,
+            ntree,
+            initkw,
+        )
+        final_state, burnin_trace, main_trace = self._run_mcmc(
+            mcmc_state, ndpost, nskip, keepevery, printevery, seed
+        )
 
         sigma = self._extract_sigma(main_trace, scale)
         first_sigma = self._extract_sigma(burnin_trace, scale)
@@ -239,7 +262,7 @@ class gbart:
 
         Parameters
         ----------
-        x_test : array (m, p) or DataFrame
+        x_test : array (p, m) or DataFrame
             The test predictors.
 
         Returns
@@ -285,7 +308,9 @@ class gbart:
         assert get_length(x1) == get_length(x2)
 
     @staticmethod
-    def _process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda, offset):
+    def _process_noise_variance_settings(
+        x_train, y_train, sigest, sigdf, sigquant, lamda, offset
+    ):
         if lamda is not None:
             return lamda, None
         else:
@@ -298,7 +323,7 @@ class gbart:
             else:
                 x_centered = x_train.T - x_train.mean(axis=1)
                 y_centered = y_train - y_train.mean()
-                    # centering is equivalent to adding an intercept column
+                # centering is equivalent to adding an intercept column
                 _, chisq, rank, _ = jnp.linalg.lstsq(x_centered, y_centered)
                 chisq = chisq.squeeze(0)
                 dof = len(y_train) - rank
@@ -336,11 +361,25 @@ class gbart:
         return prepcovars.bin_predictors(x, splits)
 
     @staticmethod
-    def _transform_input(y, offset, scale):
-        return (y - offset) / scale
+    def _transform_input(y, lamda, offset, scale):
+        y = (y - offset) / scale
+        lamda = lamda / (scale * scale)
+        return y, lamda
 
     @staticmethod
-    def _setup_mcmc(x_train, y_train, max_split, lamda, sigdf, power, base, maxdepth, ntree, initkw):
+    def _setup_mcmc(
+        x_train,
+        y_train,
+        w,
+        max_split,
+        lamda,
+        sigdf,
+        power,
+        base,
+        maxdepth,
+        ntree,
+        initkw,
+    ):
         depth = jnp.arange(maxdepth - 1)
         p_nonterminal = base / (1 + depth).astype(float) ** power
         sigma2_alpha = sigdf / 2
@@ -348,6 +387,7 @@ class gbart:
         kw = dict(
             X=x_train,
             y=y_train,
+            error_scale=w,
             max_split=max_split,
             num_trees=ntree,
             p_nonterminal=p_nonterminal,
@@ -355,17 +395,20 @@ class gbart:
             sigma2_beta=sigma2_beta,
             min_points_per_leaf=5,
         )
-        kw.update(initkw)
+        if initkw is not None:
+            kw.update(initkw)
         return mcmcstep.init(**kw)
 
     @staticmethod
     def _run_mcmc(mcmc_state, ndpost, nskip, keepevery, printevery, seed):
-        if isinstance(seed, jax.Array) and jnp.issubdtype(seed.dtype, jax.dtypes.prng_key):
+        if isinstance(seed, jax.Array) and jnp.issubdtype(
+            seed.dtype, jax.dtypes.prng_key
+        ):
             key = seed
         else:
             key = jax.random.key(seed)
         callback = mcmcloop.make_simple_print_callback(printevery)
-        return mcmcloop.run_mcmc(mcmc_state, nskip, ndpost, keepevery, callback, key)
+        return mcmcloop.run_mcmc(key, mcmc_state, nskip, ndpost, keepevery, callback)
 
     @staticmethod
     def _predict(trace, x):
@@ -379,9 +422,9 @@ class gbart:
     def _extract_sigma(trace, scale):
         return scale * jnp.sqrt(trace['sigma2'])
 
-
     def _show_tree(self, i_sample, i_tree, print_all=False):
         from . import debug
+
         trace = self._main_trace
         leaf_tree = trace['leaf_trees'][i_sample, i_tree]
         var_tree = trace['var_trees'][i_sample, i_tree]
@@ -396,7 +439,9 @@ class gbart:
         else:
             resid = bart['resid']
             alpha = bart['sigma2_alpha'] + resid.size / 2
-            norm2 = jnp.dot(resid, resid, preferred_element_type=bart['sigma2_beta'].dtype)
+            norm2 = jnp.dot(
+                resid, resid, preferred_element_type=bart['sigma2_beta'].dtype
+            )
             beta = bart['sigma2_beta'] + norm2 / 2
         sigma2 = beta / alpha
         return jnp.sqrt(sigma2) * self.scale
@@ -404,22 +449,32 @@ class gbart:
     def _compare_resid(self):
         bart = self._mcmc_state
         resid1 = bart['resid']
-        yhat = grove.evaluate_forest(bart['X'], bart['leaf_trees'], bart['var_trees'], bart['split_trees'], jnp.float32)
+        yhat = grove.evaluate_forest(
+            bart['X'],
+            bart['leaf_trees'],
+            bart['var_trees'],
+            bart['split_trees'],
+            jnp.float32,
+        )
         resid2 = bart['y'] - yhat
         return resid1, resid2
 
     def _avg_acc(self):
         trace = self._main_trace
+
         def acc(prefix):
             acc = trace[f'{prefix}_acc_count']
             prop = trace[f'{prefix}_prop_count']
             return acc.sum() / prop.sum()
+
         return acc('grow'), acc('prune')
 
     def _avg_prop(self):
         trace = self._main_trace
+
         def prop(prefix):
             return trace[f'{prefix}_prop_count'].sum()
+
         pgrow = prop('grow')
         pprune = prop('prune')
         total = pgrow + pprune
@@ -432,16 +487,21 @@ class gbart:
 
     def _depth_distr(self):
         from . import debug
+
         trace = self._main_trace
         split_trees = trace['split_trees']
         return debug.trace_depth_distr(split_trees)
 
     def _points_per_leaf_distr(self):
         from . import debug
-        return debug.trace_points_per_leaf_distr(self._main_trace, self._mcmc_state['X'])
+
+        return debug.trace_points_per_leaf_distr(
+            self._main_trace, self._mcmc_state['X']
+        )
 
     def _check_trees(self):
         from . import debug
+
         return debug.check_trace(self._main_trace, self._mcmc_state)
 
     def _tree_goes_bad(self):

{bartz-0.4.1 → bartz-0.5.0}/src/bartz/__init__.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/__init__.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -28,13 +28,5 @@ Super-fast BART (Bayesian Additive Regression Trees) in Python
 See the manual at https://gattocrucco.github.io/bartz/docs
 """
 
-from ._version import __version__
-
-from . import BART
-
-from . import debug
-from . import grove
-from . import mcmcstep
-from . import mcmcloop
-from . import prepcovars
-from . import jaxext
+from . import BART, debug, grove, jaxext, mcmcloop, mcmcstep, prepcovars  # noqa: F401
+from ._version import __version__  # noqa: F401

bartz-0.5.0/src/bartz/_version.py

@@ -0,0 +1 @@
+__version__ = '0.5.0'

{bartz-0.4.1 → bartz-0.5.0}/src/bartz/debug.py

@@ -1,21 +1,19 @@
 import functools
 
 import jax
-from jax import numpy as jnp
 from jax import lax
+from jax import numpy as jnp
 
-from . import grove
-from . import mcmcstep
-from . import jaxext
+from . import grove, jaxext
 
-def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
 
+def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
     tee = '├──'
     corner = '└──'
     join = '│  '
     space = '   '
     down = '┐'
-    bottom = '╢' # '┨' # 
+    bottom = '╢'  # '┨' #
 
     def traverse_tree(index, depth, indent, first_indent, next_indent, unused):
         if index >= len(leaf_tree):
@@ -58,7 +56,7 @@ def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
 
         indent += next_indent
         unused = unused or is_leaf
-        
+
         if unused and not print_all:
             return
 
@@ -67,58 +65,80 @@ def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
 
     traverse_tree(1, 0, '', '', '', False)
 
+
 def tree_actual_depth(split_tree):
     is_leaf = grove.is_actual_leaf(split_tree, add_bottom_level=True)
     depth = grove.tree_depths(is_leaf.size)
     depth = jnp.where(is_leaf, depth, 0)
     return jnp.max(depth)
 
+
 def forest_depth_distr(split_trees):
     depth = grove.tree_depth(split_trees) + 1
     depths = jax.vmap(tree_actual_depth)(split_trees)
     return jnp.bincount(depths, length=depth)
 
+
 def trace_depth_distr(split_trees_trace):
     return jax.vmap(forest_depth_distr)(split_trees_trace)
 
+
 def points_per_leaf_distr(var_tree, split_tree, X):
     traverse_tree = jax.vmap(grove.traverse_tree, in_axes=(1, None, None))
     indices = traverse_tree(X, var_tree, split_tree)
-    count_tree = jnp.zeros(2 * split_tree.size, dtype=jaxext.minimal_unsigned_dtype(indices.size))
+    count_tree = jnp.zeros(
+        2 * split_tree.size, dtype=jaxext.minimal_unsigned_dtype(indices.size)
+    )
     count_tree = count_tree.at[indices].add(1)
     is_leaf = grove.is_actual_leaf(split_tree, add_bottom_level=True).view(jnp.uint8)
     return jnp.bincount(count_tree, is_leaf, length=X.shape[1] + 1)
 
+
 def forest_points_per_leaf_distr(bart, X):
     distr = jnp.zeros(X.shape[1] + 1, int)
     trees = bart['var_trees'], bart['split_trees']
+
     def loop(distr, tree):
         return distr + points_per_leaf_distr(*tree, X), None
+
     distr, _ = lax.scan(loop, distr, trees)
     return distr
 
+
 def trace_points_per_leaf_distr(bart, X):
     def loop(_, bart):
         return None, forest_points_per_leaf_distr(bart, X)
+
     _, distr = lax.scan(loop, None, bart)
     return distr
 
+
 def check_types(leaf_tree, var_tree, split_tree, max_split):
     expected_var_dtype = jaxext.minimal_unsigned_dtype(max_split.size - 1)
     expected_split_dtype = max_split.dtype
-    return var_tree.dtype == expected_var_dtype and split_tree.dtype == expected_split_dtype
+    return (
+        var_tree.dtype == expected_var_dtype
+        and split_tree.dtype == expected_split_dtype
+    )
+
 
 def check_sizes(leaf_tree, var_tree, split_tree, max_split):
     return leaf_tree.size == 2 * var_tree.size == 2 * split_tree.size
 
+
 def check_unused_node(leaf_tree, var_tree, split_tree, max_split):
     return (var_tree[0] == 0) & (split_tree[0] == 0)
 
+
 def check_leaf_values(leaf_tree, var_tree, split_tree, max_split):
     return jnp.all(jnp.isfinite(leaf_tree))
 
+
 def check_stray_nodes(leaf_tree, var_tree, split_tree, max_split):
-    index = jnp.arange(2 * split_tree.size, dtype=jaxext.minimal_unsigned_dtype(2 * split_tree.size - 1))
+    index = jnp.arange(
+        2 * split_tree.size,
+        dtype=jaxext.minimal_unsigned_dtype(2 * split_tree.size - 1),
+    )
     parent_index = index >> 1
     is_not_leaf = split_tree.at[index].get(mode='fill', fill_value=0) != 0
     parent_is_leaf = split_tree[parent_index] == 0
@@ -126,6 +146,7 @@ def check_stray_nodes(leaf_tree, var_tree, split_tree, max_split):
     stray = stray.at[1].set(False)
     return ~jnp.any(stray)
 
+
 check_functions = [
     check_types,
     check_sizes,
@@ -134,6 +155,7 @@ check_functions = [
     check_stray_nodes,
 ]
 
+
 def check_tree(leaf_tree, var_tree, split_tree, max_split):
     error_type = jaxext.minimal_unsigned_dtype(2 ** len(check_functions) - 1)
     error = error_type(0)
@@ -144,15 +166,19 @@ def check_tree(leaf_tree, var_tree, split_tree, max_split):
         error |= bit
     return error
 
+
 def describe_error(error):
-    return [
-        func.__name__
-        for i, func in enumerate(check_functions)
-        if error & (1 << i)
-    ]
+    return [func.__name__ for i, func in enumerate(check_functions) if error & (1 << i)]
+
 
 check_forest = jax.vmap(check_tree, in_axes=(0, 0, 0, None))
 
+
 @functools.partial(jax.vmap, in_axes=(0, None))
 def check_trace(trace, state):
-    return check_forest(trace['leaf_trees'], trace['var_trees'], trace['split_trees'], state['max_split'])
+    return check_forest(
+        trace['leaf_trees'],
+        trace['var_trees'],
+        trace['split_trees'],
+        state['max_split'],
+    )