Added project_bite_tpu_unittests and associated config to test axlear…

[asall]

…n framework with v6e TPUs

• Currently only runs tests in axlearn/common/flash_attention/

Description

Added code to automated test runs of some of the unit tests in the Axlearn framework on a v6e TPU VM to help diagnose any issues arising from changes in Jax, TPU framework etc.

Tests

Please describe the tests that you ran on Cloud VM to verify changes.

On a v6e TPU VM, checkout axlearn (https://github.com/apple/axlearn ), configure the dependencies, and run:
pytest --no-header -v axlearn/common/flash_attention/
from within the axlearn repo.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run one-shot tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.