jaxpruner

Repair a bug, which causes Adam statue not correctly reset

Repair a bug of tree_map None

Updating scaling laws colab

Add Wanda implementation to jaxpruner (https://arxiv.org/pdf/2306.11695.pdf)

added example scaling law notebook and data