BART vectoriZed
An implementation of Bayesian Additive Regression Trees (BART) in JAX.
If you don't know what BART is, but know XGBoost, consider BART as a sort of Bayesian XGBoost. bartz makes BART run fast enough to be used in place of XGBoost/lightgbm/etc. even on large datasets.
BART is a nonparametric Bayesian regression technique. Given training 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) if n > 20,000, but using 1/20 of the memory. On GPU, the speed premium depends on sample size; it is convenient over CPU only for n > 10,000. The maximum speedup is currently 200x, on an Nvidia A100 and with at least 2,000,000 observations.
This Colab notebook runs bartz with n = 100,000 observations, p = 1000 predictors, 10,000 trees, for 2000 MCMC iterations, in 10 minutes.
BART is a very flexible method with many variants. This implementation provides only a small subset of the possible features. If you need a feature from another BART implementation or from the BART literature, please open an issue on github.
Links
- Documentation (latest release)
- Documentation (development version)
- Repository
- Code coverage
- Benchmarks
- List of BART packages
Citing bartz
Article: Petrillo (2024), "Very fast Bayesian Additive Regression Trees on GPU", arXiv:2410.23244.
To cite the software directly, including the specific version, use zenodo.
1 maintainer
Metadata
1 maintainer
Release files for bartz 0.12.1
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| bartz-0.12.1.tar.gz | 181.2 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| bartz-0.12.1-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 404.3 kB
Release files / bartz-0.12.1.tar.gz
| Download URL | bartz-0.12.1.tar.gz |
|---|---|
| Size | 181.2 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
6cebefa1d146790c5a4cbcb1def63f510cab4f467089113b8b449425d21de802
|
|
BLAKE2b-256 checksum How to use checksums |
1bb09c86a225cbf6ecd8f60cf61574997a43e0741b2deb860cb42195cb4ef98e
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.12.5 {"installer":{"name":"uv","version":"0.12.5","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
|
Release files / bartz-0.12.1-py3-none-any.whl
| Download URL | bartz-0.12.1-py3-none-any.whl |
|---|---|
| Size | 223.1 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
7abf99e55c1082245500f2ccd484c1bc3a9e542c3eab680852b964d702aded71
|
|
BLAKE2b-256 checksum How to use checksums |
6bf686e4e5bda0866f41cf38df39ecd89b073fa87a499cac2de709c2e95ba57f
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.12.5 {"installer":{"name":"uv","version":"0.12.5","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
|
