Package: jax-relax
ReLaxOverview|Installation|Tutorials|Documentation|Citing ReLaxOverviewReLax(Recourse ExplanationLibrary in Jax) is an
efficient and scalable benchmarking library for recourse and
counterfactual explanations, built on top ofjax. By leveraging language
primitives such asvectorization,parallelization, andjust-in-timecompilation injax,ReLaxoffers massive
speed improvements in generating individual (or local) explanations for
predictions made by Machine Learning algorithms.Some of the key features are as follows:🏃Fast and scalablerecourse generation.🚀Acceleratedovercpu,gpu,tpu.🪓Comprehensiveset of recourse methods implemented for
benchmarking.👐CustomizableAPI to enable the building of entire modeling and
interpretation pipelines for new recourse algorithms.Installationpipinstalljax-relax# Or install the latest version of `jax-relax`pipinstallgit+https://github.com/BirkhoffG/jax-relax.gitTo futher unleash the power of accelerators (i.e., GPU/TPU), we suggest
to first install this library viapip install jax-relax. Then, follow
steps in theofficial install
guidelinesto install the
right version for GPU or TPU.Dive intoReLaxReLaxis a recourse explanation library for explaining (any) JAX-based
ML models. We believe that it is important to give users flexibility to
choose how to useReLax. You canonly use methods implemeted inReLax(as a recourse methods
library);build a pipeline usingReLaxto define data module, training ML
models, and generating CF explanation (for constructing recourse
benchmarking pipeline).ReLaxas a Recourse Explanation LibraryWe introduce basic use cases of using methods inReLaxto generate
recourse explanations. For more advanced usages of methods inReLax,
See thistutorials.fromrelax.methodsimportVanillaCFfromrelaximportDataModule,MLModule,generate_cf_explanations,benchmark_cfsfromsklearn.datasetsimportmake_classificationfromsklearn.model_selectionimporttrain_test_splitimportfunctoolsasftimportjaxLet’s first generate synthetic data:xs,ys=make_classification(n_samples=1000,n_features=10,random_state=42)train_xs,test_xs,train_ys,test_ys=train_test_split(xs,ys,random_state=42)Next, we fit an MLP model for this data. Note that this model can be any
model implmented in JAX. We will use theMLModuleinReLaxas an example.model=MLModule()model.train((train_xs,train_ys),epochs=10,batch_size=64)Generating recourse explanations are straightforward. We can simply callgenerate_cfof an implemented recourse method to generateonerecourse explanation:vcf=VanillaCF(config={'n_steps':1000,'lr':0.05})cf=vcf.generate_cf(test_xs[0],model.pred_fn)assertcf.shape==test_xs[0].shapeOr generate a bunch of recourse explanations withjax.vmap:generate_fn=ft.partial(vcf.generate_cf,pred_fn=model.pred_fn)cfs=jax.vmap(generate_fn)(test_xs)assertcfs.shape==test_xs.shapeReLaxfor Building Recourse Explanation PipelinesThe above example illustrates the usage of the decoupledrelax.methodsto generate recourse explanations. However, users are required to write
boilerplate code for tasks such as data preprocessing, model training,
and generating recourse explanations with feature constraints.ReLaxadditionally offers a one-liner framework, streamlining the
process and helping users in building a standardized pipeline for
generating recourse explanations. You can write three lines of code to
benchmark recourse explanations:data_module=DataModule.from_numpy(xs,ys)exps=generate_cf_explanations(vcf,data_module,model.pred_fn)benchmark_cfs([exps])SeeGetting Started with
ReLaxfor an end-to-end example of usingReLax.Supported Recourse MethodsReLaxcurrently provides implementations of 9 recourse explanation
methods.MethodTypePaper TitleRefVanillaCFNon-ParametricCounterfactual Explanations without Opening the Black Box: Automated Decisions and the GDPR.[1]DiverseCFNon-ParametricExplaining Machine Learning Classifiers through Diverse Counterfactual Explanations.[2]ProtoCFSemi-ParametricInterpretable Counterfactual Explanations Guided by Prototypes.[3]CounterNetParametricCounterNet: End-to-End Training of Prediction Aware Counterfactual Explanations.[4]GrowingSphereNon-ParametricInverse Classification for Comparison-based Interpretability in Machine Learning.[5]CCHVAESemi-ParametricLearning Model-Agnostic Counterfactual Explanations for Tabular Data.[6]VAECFParametricPreserving Causal Constraints in Counterfactual Explanations for Machine Learning Classifiers.[7]CLUESemi-ParametricGetting a CLUE: A Method for Explaining Uncertainty Estimates.[8]L2CParametricFeature-based Learning for Diverse and Privacy-Preserving Counterfactual Explanations[9]CitingReLaxTo cite this repository:@software{relax2023github,
  author ={Hangzhi Guo and Xinchang Xiong and Amulya Yadav},
  title ={{R}e{L}ax: Recourse Explanation Library in Jax},
  url ={http://github.com/birkhoffg/jax-relax},
  version ={0.2.0},
  year ={2023},}