p.enthalabs

GitHub - google-research/kauldron: Modular, scalable library to train ML models

Source: https://github.com/google-research/kauldron

Skip to content

Navigation Menu

[](https://github.com/)

Sign in

* Products

* Platform

- ![Image 1!Image 2GitHub Repository hosting, code review, and collaboration](https://github.com/platform)

- ![Image 3!Image 4Copilot AI coding agent for tasks, projects and automated work](https://github.com/features/copilot)

* Build & Ship

- Actions CI/CD and automation

- Code Quality Quality analysis, fixes, and enforcement

- Codespaces Cloud development environments

* Secure

- Code Security Vulnerability detection and remediation

- Secret Protection Secret scanning and leak prevention

All Features

* Developers

* Learn

- Documentation

- Training

- Certifications

* Build

- Build with Copilot

- Customize Copilot

- Integrations

* Explore

- Trending repos

- GitHub Sponsors

- Events

* Changelog

- ![Image 5!Image 6Changelog Your source for the latest features, improvements, and fixes across GitHub products, all in one place.](https://github.blog/changelog/)

* Businesses

* By Stage

- Enterprises

- Team

- Startups

- Nonprofits

* By Industry

- Healthcare

- Financial services

- Manufacturing

- Government

- See all industries

* Evaluate GitHub

- Why GitHub

- Customer Stories

- Business insights

- Trust Center

* Add-ons

- GitHub Advanced Security Enterprise-grade code security

- Copilot for business Enterprise-grade AI features

- Premium Support Enterprise-grade 24/7 support

* Resources

* Learn & Connect

- Documentation

- Resource center

- Blog

- Events & webinars

- Community Forum

- Customer Stories

- Integrations

* Support

- Customer support

- Training

- GitHub Status

* Company

- About

- Partners

- Changelog

* GitHub Universe

- ![Image 7!Image 8GitHub Universe The event for developers, by developers — talks, workshops, and what’s next for GitHub.](https://githubuniverse.com/)

- Pricing

Search/

Sign in

Sign up

Appearance settings

You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert

{{ message }}

1. google-research 2. kauldron

Repository navigation

- Code

- Issues 7(7)

- Pull requests 88(88)

- Actions

- Security and quality

- Insights

More items

[](https://github.com/google-research/kauldron)

!Image 9: google-research

**kauldron**

Public

- NotificationsYou must be signed in to change notification settings

- Fork 36(36)

- Star 352(352)You must be signed in to star a repository

About

Modular, scalable library to train ML models

kauldron.readthedocs.io/

Topics

deep-learningjaxpython

Resources

Readme

Apache-2.0 license

Contributing

Contributing

Activity

Custom properties

Stars

**352** stars

Watchers

**5** watching

Forks

**36** forks

Report repository

main

**83**Branches**12**Tags

[](https://github.com/google-research/kauldron/branches)[](https://github.com/google-research/kauldron/tags)

Go to file

Code

Open more actions menu

Latest commit

!Image 10: author

The kauldron Authors

Support arbitrary PyTrees of arrays in `kd.metrics.truncate_field`.

Open commit details success

Oct 5, 2026

b3b376a·Oct 5, 2026

History

1,206 Commits

Open commit details

[](https://github.com/google-research/kauldron/commits/main/)1,206 Commits

Folders and files

| Name | Name | Last commit message | Last commit date | | --- | --- | --- | --- |

| [.github/workflows](https://github.com/google-research/kauldron/tree/main/.github/workflows "This path skips through empty directories") | [.github/workflows](https://github.com/google-research/kauldron/tree/main/.github/workflows "This path skips through empty directories") | [Kauldron release 1.4.3](https://github.com/google-research/kauldron/commit/541413debfda669a3727a85ab4d54a5838157c71 "Kauldron release 1.4.3 PiperOrigin-RevId: 929762435") | Jun 10, 2026 |

| [.vscode](https://github.com/google-research/kauldron/tree/main/.vscode ".vscode") | [.vscode](https://github.com/google-research/kauldron/tree/main/.vscode ".vscode") | [Code update](https://github.com/google-research/kauldron/commit/526e050b0083e13352d4c77e7049465b52ab5b02 "Code update PiperOrigin-RevId: 547159480") | Jul 11, 2023 |

| [docs](https://github.com/google-research/kauldron/tree/main/docs "docs") | [docs](https://github.com/google-research/kauldron/tree/main/docs "docs") | [Replace CollectingState and CollectFirstState by AutoState in Kauldron](https://github.com/google-research/kauldron/commit/4c09b931e7ead8b43e9960458dd4990f1bc80f97 "Replace CollectingState and CollectFirstState by AutoState in Kauldron PiperOrigin-RevId: 941034705") | Jul 1, 2026 |

| [examples](https://github.com/google-research/kauldron/tree/main/examples "examples") | [examples](https://github.com/google-research/kauldron/tree/main/examples "examples") | [Code update](https://github.com/google-research/kauldron/commit/045874ff55de3be14e7159b39497be857c8f0ed2 "Code update PiperOrigin-RevId: 941587426") | Jul 2, 2026 |

| [kauldron](https://github.com/google-research/kauldron/tree/main/kauldron "kauldron") | [kauldron](https://github.com/google-research/kauldron/tree/main/kauldron "kauldron") | [Support arbitrary PyTrees of arrays in `kd.metrics.truncate_field`.](https://github.com/google-research/kauldron/commit/b3b376acd02fce0270d53237e071e1cf65bf1d5a "Support arbitrary PyTrees of arrays in `kd.metrics.truncate_field`. PiperOrigin-RevId: 993966716") | Oct 5, 2026 |

| [.gitignore](https://github.com/google-research/kauldron/blob/main/.gitignore ".gitignore") | [.gitignore](https://github.com/google-research/kauldron/blob/main/.gitignore ".gitignore") | [Fix open-source tests](https://github.com/google-research/kauldron/commit/307dbb4331f29fb2cec8be8d32c432af78b255e6 "Fix open-source tests PiperOrigin-RevId: 698002092") | Nov 19, 2024 |

| [.pylintrc](https://github.com/google-research/kauldron/blob/main/.pylintrc ".pylintrc") | [.pylintrc](https://github.com/google-research/kauldron/blob/main/.pylintrc ".pylintrc") | [Code update](https://github.com/google-research/kauldron/commit/526e050b0083e13352d4c77e7049465b52ab5b02 "Code update PiperOrigin-RevId: 547159480") | Jul 11, 2023 |

| [.readthedocs.yaml](https://github.com/google-research/kauldron/blob/main/.readthedocs.yaml ".readthedocs.yaml") | [.readthedocs.yaml](https://github.com/google-research/kauldron/blob/main/.readthedocs.yaml ".readthedocs.yaml") | [New attempt to fix the apitree readthedoc issue](https://github.com/google-research/kauldron/commit/57ded1fe5a10342c1eb506e1339d8757a7fa624a "New attempt to fix the apitree readthedoc issue PiperOrigin-RevId: 700640050") | Nov 27, 2024 |

| [CHANGELOG.md](https://github.com/google-research/kauldron/blob/main/CHANGELOG.md "CHANGELOG.md") | [CHANGELOG.md](https://github.com/google-research/kauldron/blob/main/CHANGELOG.md "CHANGELOG.md") | [hotfix release 1.4.4 for cuda compatibility issue](https://github.com/google-research/kauldron/commit/7da159080ac678b3346b335471f5b5c85be90bce "hotfix release 1.4.4 for cuda compatibility issue PiperOrigin-RevId: 929814140") | Jun 10, 2026 |

| [CONTRIBUTING.md](https://github.com/google-research/kauldron/blob/main/CONTRIBUTING.md "CONTRIBUTING.md") | [CONTRIBUTING.md](https://github.com/google-research/kauldron/blob/main/CONTRIBUTING.md "CONTRIBUTING.md") | [Create CONTRIBUTING.md](https://github.com/google-research/kauldron/commit/6b9e3ef9698add236ba6cc111cad1f5609b35c5b "Create CONTRIBUTING.md") | Jun 19, 2023 |

| [LICENSE](https://github.com/google-research/kauldron/blob/main/LICENSE "LICENSE") | [LICENSE](https://github.com/google-research/kauldron/blob/main/LICENSE "LICENSE") | [Code update](https://github.com/google-research/kauldron/commit/526e050b0083e13352d4c77e7049465b52ab5b02 "Code update PiperOrigin-RevId: 547159480") | Jul 11, 2023 |

| [README.md](https://github.com/google-research/kauldron/blob/main/README.md "README.md") | [README.md](https://github.com/google-research/kauldron/blob/main/README.md "README.md") | [Add a citation to the README](https://github.com/google-research/kauldron/commit/7966ee4bfad56970b4900dc0e121a717e07acd96 "Add a citation to the README PiperOrigin-RevId: 830348557") | Nov 10, 2025 |

| [pyproject.toml](https://github.com/google-research/kauldron/blob/main/pyproject.toml "pyproject.toml") | [pyproject.toml](https://github.com/google-research/kauldron/blob/main/pyproject.toml "pyproject.toml") | [Remove jax dependencies from kauldron xm launcher and kontext.](https://github.com/google-research/kauldron/commit/5aaf6135bc889adbcbbbdd643f9d332b66ee6a03 "Remove jax dependencies from kauldron xm launcher and kontext. PiperOrigin-RevId: 978508366") | Sep 9, 2026 | | View all files |

Repository files navigation

- README

- Contributing

- Apache-2.0 license

More items

kauldron

[](https://github.com/google-research/kauldron#kauldron)

![Image 11: Unittests](https://github.com/google-research/kauldron/actions/workflows/pytest_and_autopublish.yml)![Image 12: PyPI version](https://badge.fury.io/py/kauldron)![Image 13: Documentation Status](https://kauldron.readthedocs.io/en/latest/?badge=latest)

Kauldron is a library for training machine learning models, optimized for **research velocity** and **modularity**.

**Modularity**:

- All parts of Kauldron are self-contained, so can be used independently outside Kauldron.

- Use any dataset (TFDS, Grain, SeqIO, your custom pipeline), any (flax) model, any optimizer,... Kauldron provides the glue that link everything together.

- Everything can be customized and overwritten (e.g. sweep over models architecture, overwrite any inner layer parameter,...)

**Research velocity**:

- Everything should work out-of the box. The example configs can be used and customized as a starting point.

- Colab-first workflow for easy prototyping and fast iteration

- Polished user experience (integrated XM plots, profiler, post-mortem debugging on borg, runtime shape checking, and many others...). Open an issue..

Citing Kauldron

[](https://github.com/google-research/kauldron#citing-kauldron)

If Kauldron was helpful for a publication, please cite this repository:

``` @software{kauldron2025github, author = {Klaus Greff and Etienne Pot and Mehdi S. M. Sajjadi}, title = {{Kauldron}: A neural network training framework, optimized for research velocity and modularity.}, url = {https://github.com/google-research/kauldron}, version = {1.3.0}, year = {2025}, } ```

_This is not an officially supported Google product._

About

Modular, scalable library to train ML models

kauldron.readthedocs.io/

Topics

deep-learningjaxpython

Resources

Readme

Apache-2.0 license

Contributing

Contributing

Activity

Custom properties

Stars

**352** stars

Watchers

**5** watching

Forks

**36** forks

Report repository

Releases12(12)

v1.4.4 Latest Jun 10, 2026

+ 11 releases

Contributors37(37)

- ![Image 14: @Qwlouse](https://github.com/Qwlouse)

- ![Image 15: @aravindhm](https://github.com/aravindhm)

- ![Image 16: @h-joo](https://github.com/h-joo)

- ![Image 17: @oprypin](https://github.com/oprypin)

- ![Image 18: @untom](https://github.com/untom)

- ![Image 19: @hawkinsp](https://github.com/hawkinsp)

- ![Image 20: @samuelepapa](https://github.com/samuelepapa)

- ![Image 21: @aaronsarna](https://github.com/aaronsarna)

- ![Image 22: @wwkong](https://github.com/wwkong)

- ![Image 23: @pwohlhart](https://github.com/pwohlhart)

- ![Image 24: @jrruijli](https://github.com/jrruijli)

- ![Image 25: @PeterZhizhin](https://github.com/PeterZhizhin)

- ![Image 26: @tkipf](https://github.com/tkipf)

+ 23 contributors

Languages

- Python 97.5%

- Jupyter Notebook 2.5%

Footer

[](https://github.com/) © 2026 GitHub,Inc.

Footer navigation

- Terms

- Privacy

- Security

- Status

- Community

- Docs

- Contact

- Manage cookies

- Do not share my personal information

You can’t perform that action at this time.