Add llama test

[zpcore]

Resolve #6.

Note, there are two issues in torch_xla run, as in torchprime/experimental/torchax_models/test/test_llama.py

1. Error between torch_xla and native pytorch CPU run is qute obvious, at level of ~ 0.1.
2. scaled_dot_product is not been replaced with flash_attention kernel for torch_xla. This requires the code change to the model.

I kind of messed up PR #26 with rebased to a commit. Opened this new one instead...

[zpcore]

Create issue #30 to track the output numerical error.

[tengyifei]

LGTM but the PR needs to be formatted. Check out the contributing section: you can run yapf --recursive -i '*.py' torchprime launcher to format everything.