Add Tunix demo in miniGPT (#269)

* Add Tunix demo in miniGPT

Unfortunately Qwix has issues with both the new sharding annotation
and KV caching previously added; so I had to revert both

* get rid of the Tunix training widget

* get rid of TFDS widget as well

* Remove troubling metadata in miniGPT

* Address comments in miniGPT

* Add back many tags in miniGPT

* Add more tags back in miniGPT

* Attempt to fix test failure

* Another attempt to fix test

* Add Tunix as a dependency for miniGPT

* use Grain instead of TFDS in miniGPT LoRA finetuning

* Really switch to Grain in miniGPT LoRA finetuning

* remove null in execution_count to fix miniGPT test failure

* add nbval-ignore-output tags in miniGPT

* one more tag

* Add tag

* Address feedback

* Fix test error

* Add back comments

Author: windmaple

docs/requirements.txt

@@ -12,6 +12,7 @@ pandas
 penzai
 scikit-learn
 tiktoken
+google-tunix[prod]
 
 # install jax-ai-stack from current directory
 .