Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Updated TPU unittests to cover more TPU types / JAX versions #623

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Updated TPU unittests to cover more TPU types / JAX versions #623

wants to merge 3 commits into from

Conversation

asall
Copy link
Contributor

@asall asall commented Mar 18, 2025

Description

  • Expanded test to v5e, v5p and JAX 0.5.1 and nightly
  • Allow for different projects/networks for test VMs
  • Merges tests into existing project_bite_tpu_e2e DAG to minimize number of DAG entries

Tests

Tested with v5e, v5p TPU VMs

Checklist

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

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

@asall
Updated TPU unittests to cover more TPU types / JAX versions
4a5fd5d 
- Expanded test to v5e, v5p and JAX 0.5.1 and nightly
- Allow for different projects/networks for test VMs
- Merges tests into existing project_bite_tpu_e2e DAG to minimize number of DAG entries
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@richardsliu richardsliu richardsliu approved these changes

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

@parambole parambole Awaiting requested review from parambole parambole is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner

@michelle-yooh michelle-yooh Awaiting requested review from michelle-yooh michelle-yooh is a code owner

@jiya-zhang jiya-zhang Awaiting requested review from jiya-zhang jiya-zhang is a code owner

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@asall @richardsliu