Do any JAX experts know if there is an equivalent to https://captum.ai/ https://captum.ai/ - a model interpretability library for pytorch? In particular i want

Do any JAX experts know if there is an equivalent to https://captum.ai/ https://captum.ai/ - a model interpretability library for pytorch? In particular i want to be able to measure feature importance on both inputs and internal layers on a sample by sample basis. This is the only thing currently holding me back from using JAX right now. Alternatively a simle to read/understand/port implementation of DeepLIFT would work too. thanks