Skip to content

Adding basic tpu github tests#336

Merged
copybara-service[bot] merged 1 commit intomainfrom
tunix_tests
Sep 26, 2025
Merged

Adding basic tpu github tests#336
copybara-service[bot] merged 1 commit intomainfrom
tunix_tests

Conversation

@mydatascience
Copy link
Collaborator

@mydatascience mydatascience commented Sep 18, 2025

Adding github tests on tpu machines

Using current tests to run pytest on tpu machines.

Tests need to be refactored after fixes apply. (we need to run all tests after code change)

@mydatascience mydatascience changed the title [Draft] Adding basic github tests Adding basic github tests Sep 25, 2025
@mydatascience mydatascience changed the title Adding basic github tests Adding basic tpu github tests Sep 25, 2025
Copy link
Collaborator

@wang2yn84 wang2yn84 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you! Can you squash all the commits into 1 before submit?

import jax
print(f'JAX version: {jax.__version__}')
print(f'JAX devices: {jax.devices()}')
print(f'TPU available: {len(jax.devices()) > 0}')
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If it's running on CPU, this will pass too. Consider check jax.devices()[0].platform == 'tpu'

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

Signed-off-by: Vladimir Suvorov <[email protected]>

Some initial tests

Signed-off-by: Vladimir Suvorov <[email protected]>

Some initial tests

Signed-off-by: Vladimir Suvorov <[email protected]>

Some initial tests

Signed-off-by: Vladimir Suvorov <[email protected]>

Some initial tests

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

TPU simple tests

Signed-off-by: Vladimir Suvorov <[email protected]>

Refactored

Signed-off-by: Vladimir Suvorov <[email protected]>

Refactored

Signed-off-by: Vladimir Suvorov <[email protected]>
Copy link
Collaborator

@wang2yn84 wang2yn84 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you very much!

@copybara-service copybara-service bot merged commit e016771 into main Sep 26, 2025
7 checks passed
Tianjiao-Yu pushed a commit to PLAN-Lab/tunix that referenced this pull request Oct 5, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants