In this proposed project we build JaxModel, which provides a high-level API for training and deploying Machine Learning Models using JAX framework. The main goal for JaxModel will be to provide an API for building, training, and testing deep learning models which are compatible with the existing DeepChem code base which already provides features like Molnet, Featurisers, Splitters, Transformers, and other deep learning infrastructures. During the period of GSOC, I will also be adding a notebook tutorial specifying the usage of JaxModel. If time permits, In the same tutorial I will also use an advanced JAX-based model preferably (Graph-based model or Physics Informed Neural Networks). This will be the beginning of implementing JAX-based models into the main DeepChem repo.



M. Vignesh Venkataraman


  • Bharath Ramsundar
  • Nathan