Archive RPAv2 and Delete Torchax-pt path.

[QiliangCui]

Description

1. Archive RPAvv2
2. Delete Torchax-pt path but leave some Pallas attention files.

FIXES: b/438779127

Tests

Waiting for
https://buildkite.com/tpu-commons/tpu-commons-ci/builds/2396

[github-actions]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[bythew3i]

Since we torchax is fixed.. Can we update the tests that were temporarily skipped? (maybe in a separated PR)

• https://screenshot.googleplex.com/3QLdS7R4xACc7om
• https://screenshot.googleplex.com/4iJugJAto83XpMz

[kyuyeunk]

Any reason for deleting this file? I have a draft PR that relies on this code path: #512

[QiliangCui]

Because this file is for torchax-pt path which will not be used anymore.

[QiliangCui]

Since we torchax is fixed.. Can we update the tests that were temporarily skipped? (maybe in a separated PR)

hm... the tests are still failing with some type errors like

TypeError: Cannot interpret 'torch.bfloat16' as a data type

AttributeError: 'NoneType' object has no attribute 'dtype'

Should be able to be fixed.
It may take me some time...

I will address in anther PR to avoid making this one too big.

Is it ok?

[kyuyeunk]

I'm planning to use this file in my torchax-jax refactoring. Can you remove deletion of this file from this PR?

[QiliangCui]

hm.. why?
I think this file is used only for torchax-pt path.

If this is not deleted, neither can tpu_commons.models.torchax.torchax_wrapper.

Could you try which part of the file is useful for you and you check it it?

[QiliangCui]

@lsy323 , what do you think?

[kyuyeunk]

Synced offline. Concluded that this file will not be removed since refactoring in #512 uses this file for torchax-jax path in favor of custom JAX attention layer. Also, other files (like tpu_commons.models.torchax.torchax_wrapper) can be deleted without any problem.

[kyuyeunk]

Hi @QiliangCui, as discussed, can you omit tpu_commons/attention/backends/pallas_torchax.py from this PR and submit this? As mentioned in #512, it will make the refactoring PR much simpler once torchax-pt related code has been removed.