Do any JAX experts know if there is an equivalent to 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