Run tests on TPU #21425

kharshith-k · 2025-06-26T07:04:44Z

Running tests on V6 TPU runner.

… configs

… hosted TPU based runner

updated build file path

Updated tpu_build job of actions.yml with specific runner label

Added container section

codecov-commenter · 2025-06-26T07:10:25Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.48%. Comparing base (6d06085) to head (4b84da1).
⚠️ Report is 18 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21425      +/-   ##
==========================================
- Coverage   82.66%   82.48%   -0.19%     
==========================================
  Files         577      577              
  Lines       59419    59506      +87     
  Branches     9313     9330      +17     
==========================================
- Hits        49121    49084      -37     
- Misses       7898     8010     +112     
- Partials     2400     2412      +12

Flag	Coverage Δ
keras	`82.30% <ø> (-0.19%)`	⬇️
keras-jax	`62.89% <ø> (-0.44%)`	⬇️
keras-numpy	`57.55% <ø> (-0.02%)`	⬇️
keras-openvino	`34.35% <ø> (+<0.01%)`	⬆️
keras-tensorflow	`64.12% <ø> (-0.01%)`	⬇️
keras-torch	`63.61% <ø> (-0.02%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…tpu.txt

Progress bar would always report the starting batch + 1 at the end of the batch. Now it takes into account `steps_per_execution` for the last batch reported. Fixes keras-team#20861

Using `keras.ops.math.logsumexp` with an int for `axis` in a functional model would throw an error.

…eras-team#21429) Arbitrary functions and classes are not allowed. - Made `Operation` extend `KerasSaveable`, this required moving imports to avoid circular imports - `Layer` no longer need to extend `KerasSaveable` directly - Made feature space `Cross` and `Feature` extend `KerasSaveable` - Also dissallow public function `enable_unsafe_deserialization`

…developed dtypes_new_test.py to use requires_tpu marker

hertschuh

Thanks for the updates!

requirements-tensorflow-tpu.txt

.github/workflows/tpu-tests-jax.yml

kharshith-k · 2025-11-06T13:01:32Z

An observed was made that tests were silently falling back to CPU during execution even when TPU was available while running them on a TPU VM which lead to the suspicion that the same could be happening with current passing tests. its a good idea to log the device type on which tests are getting executed (cpu/cuda/tpu).

hertschuh

Oh, one last thing.

We want these tests to run only after the PR is approved, like the kokoro GPU tests. Do you know how to do that?

kharshith-k · 2025-11-07T09:16:42Z

Hi @hertschuh, I'm not sure how to do that exactly, Should I add the TPU specific workflow in .kokoro/github/ubuntu/gpu/build.sh? Can you please share some resources regarding the same if you have any? In the mean time I'll do some reading from my end too!

hertschuh · 2025-11-07T21:03:56Z

Hi @hertschuh, I'm not sure how to do that exactly, Should I add the TPU specific workflow in .kokoro/github/ubuntu/gpu/build.sh? Can you please share some resources regarding the same if you have any? In the mean time I'll do some reading from my end too!

It looks like this would be a way to do it: https://github.com/orgs/community/discussions/25372#discussioncomment-3247688

But it means you need a separate yml file.

…the PR review is approved

kharshith-k · 2025-11-11T10:00:52Z

Hello @hertschuh , I've added the condition to check if the PR is reviewed and approved before running the TPU tests job. Can you please let me know if the change is valid?

hertschuh

Thanks!

.github/workflows/actions.yml

hertschuh

Last 2 nitpicks

.github/workflows/tpu_tests.yml

… workflow

kharshith-k added 17 commits June 16, 2025 15:23

added requirements-tensorflow-tpu.txt and tpu configuration in .kokoro

1d7c685

updated .kokoro/github/ubuntu/tpu/build.sh with jax and torch backend…

19b5e6b

… configs

Merge branch 'keras-team:master' into tf-tpu

d203ca3

Changed the tpu CI config files path to .github from .kokoro

f45e5d0

Added new job in .github/workflows/actions.yml to run TPU tests

6771cc0

fixed runs-on option in acvtions.yml for tpu_build job to run on self…

87d36e7

… hosted TPU based runner

Added another runner in the actions TPU job

9901298

Update continuous.cfg

be97210

updated build file path

Update presubmit.cfg

a1cd5c3

updated build file path

Merge branch 'keras-team:master' into tf-tpu

c5e3a5c

Update actions.yml

f0ab676

Updated tpu_build job of actions.yml with specific runner label

Developed Dockerfile for TPU build job in actions.yml

09161d7

Merge branch 'keras-team:master' into tf-tpu

9a3948f

Update actions.yml

058fdff

Added container section

Included few more runners in tpu_build job

d47e39e

Merge branch 'keras-team:master' into tf-tpu

a6a59d7

Using linux-x86-ct6e-44-1tpu

ba4f6ae

google-ml-butler bot added the size:M label Jun 26, 2025

google-ml-butler bot assigned gbaned Jun 26, 2025

kharshith-k and others added 10 commits June 30, 2025 08:34

Modified requirement-commmon.txt and updated requirements-tensorflow-…

a5a3624

…tpu.txt

Added Dtypes_TPU_tests.py and requirements-jax-tpu.txt

b9998af

Progress bar now handles steps_per_execution. (keras-team#21422)

f68be97

Progress bar would always report the starting batch + 1 at the end of the batch. Now it takes into account `steps_per_execution` for the last batch reported. Fixes keras-team#20861

Fix symbolic call of logsumexp with int axis. (keras-team#21428)

1018abf

Using `keras.ops.math.logsumexp` with an int for `axis` in a functional model would throw an error.

commented tensorflow deps

cb639c5

Added log of dtypes_test_tpu.py and the test script for the same

c0d1743

modified dtypes_test_tpu.py as per pre-commit standards

306e6e7

Added TPU initiaization and teardown functionalities in conftest.py, …

4e584fc

…developed dtypes_new_test.py to use requires_tpu marker

Added dtypes_test_TPU.py and dtypes_new_test.py, modified conftest.py

bb09e95

Updated the workflow and job names for TPU

4e0d3af

kharshith-k changed the title ~~Running TPU tests on linux-x86-ct6e-44-1tpu~~ Run tests on TPU Nov 5, 2025

hertschuh reviewed Nov 5, 2025

View reviewed changes

requirements-tensorflow-tpu.txt Outdated Show resolved Hide resolved

requirements-tensorflow-tpu.txt Outdated Show resolved Hide resolved

.github/workflows/tpu-tests-jax.yml Show resolved Hide resolved

.github/workflows/tpu-tests-jax.yml Outdated Show resolved Hide resolved

Added TPU job in actions.yml

0d0b009

Added TPU config print line

65252a4

hertschuh reviewed Nov 6, 2025

View reviewed changes

Added condition to TPU test job so that it gets triggered only after …

1769e32

…the PR review is approved

hertschuh approved these changes Nov 11, 2025

View reviewed changes

google-ml-butler bot added kokoro:force-run ready to pull Ready to be merged into the codebase labels Nov 11, 2025

kokoro-team removed the kokoro:force-run label Nov 11, 2025

hertschuh requested changes Nov 11, 2025

View reviewed changes

.github/workflows/actions.yml Outdated Show resolved Hide resolved

Added TPU specific tests in seperate workflow with PR approval condition

c44ce18

google-ml-butler bot removed the ready to pull Ready to be merged into the codebase label Nov 12, 2025

hertschuh approved these changes Nov 12, 2025

View reviewed changes

google-ml-butler bot added kokoro:force-run ready to pull Ready to be merged into the codebase labels Nov 12, 2025

kokoro-team removed the kokoro:force-run label Nov 12, 2025

hertschuh reviewed Nov 12, 2025

View reviewed changes

.github/workflows/tpu_tests.yml Outdated Show resolved Hide resolved

.github/workflows/tpu_tests.yml Outdated Show resolved Hide resolved

Removed pull_request condition for workflow execution and renamed the…

4b84da1

… workflow

google-ml-butler bot removed the ready to pull Ready to be merged into the codebase label Nov 13, 2025

hertschuh approved these changes Nov 13, 2025

View reviewed changes

google-ml-butler bot added kokoro:force-run ready to pull Ready to be merged into the codebase labels Nov 13, 2025

kokoro-team removed the kokoro:force-run label Nov 13, 2025

hertschuh merged commit 2730d52 into keras-team:master Nov 13, 2025
12 checks passed

Run tests on TPU #21425

Run tests on TPU #21425

Conversation

kharshith-k commented Jun 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov-commenter commented Jun 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kharshith-k commented Nov 6, 2025

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

kharshith-k commented Nov 7, 2025

Uh oh!

hertschuh commented Nov 7, 2025

Uh oh!

kharshith-k commented Nov 11, 2025

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

kharshith-k commented Jun 26, 2025 •

edited

Loading

codecov-commenter commented Jun 26, 2025 •

edited

Loading