Gattocrucco / bartz

A JAX implementation of BART (Bayesian Additive Regression Trees)

Home Page:https://gattocrucco.github.io/bartz/docs

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

PyPI

BART vectoriZed

A branchless vectorized implementation of Bayesian Additive Regression Trees (BART) in JAX.

BART is a nonparametric Bayesian regression technique. Given predictors $X$ and responses $y$, BART finds a function to predict $y$ given $X$. The result of the inference is a sample of possible functions, representing the uncertainty over the determination of the function.

This Python module provides an implementation of BART that runs on GPU, to process large datasets faster. It is also good on CPU. Most other implementations of BART are for R, and run on CPU only.

On CPU, bartz runs at the speed of dbarts (the fastest implementation I know of), but using half the memory. On GPU, the speed premium depends on sample size; with 50000 datapoints and 5000 trees, on an Nvidia Tesla V100 GPU it's 12 times faster than a single Apple M1 CPU core, and this factor is linearly proportional to the number of datapoints.

The maximum practically realizable speedup is currently 200x.

Links

Other BART packages

About

A JAX implementation of BART (Bayesian Additive Regression Trees)

https://gattocrucco.github.io/bartz/docs

License:MIT License


Languages

Language:Python 97.4%Language:Makefile 2.2%Language:HTML 0.4%