{"id":42102710,"url":"https://github.com/google/torchax","last_synced_at":"2026-07-24T05:01:01.915Z","repository":{"id":317927602,"uuid":"1069392963","full_name":"google/torchax","owner":"google","description":"torchax is a PyTorch frontend for JAX. It gives JAX the ability to author JAX programs using familiar PyTorch syntax. It also provides JAX-Pytorch interoperability, meaning, one can mix JAX \u0026 Pytorch syntax together when authoring ML programs, and run it in every hardware JAX can run.","archived":false,"fork":false,"pushed_at":"2026-07-03T00:31:51.000Z","size":1588,"stargazers_count":231,"open_issues_count":19,"forks_count":36,"subscribers_count":5,"default_branch":"main","last_synced_at":"2026-07-24T03:06:53.719Z","etag":null,"topics":["ai","interoperability","jax","machine-learning","pytorch","tpu"],"latest_commit_sha":null,"homepage":"https://google.github.io/torchax/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/google.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-10-03T21:32:18.000Z","updated_at":"2026-07-21T09:44:09.000Z","dependencies_parsed_at":"2025-10-04T00:14:59.230Z","dependency_job_id":"ebff87c0-39d4-4bfb-9d2e-352e633b4201","html_url":"https://github.com/google/torchax","commit_stats":null,"previous_names":["google/torchax"],"tags_count":4,"template":false,"template_full_name":null,"purl":"pkg:github/google/torchax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google%2Ftorchax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google%2Ftorchax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google%2Ftorchax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google%2Ftorchax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/google","download_url":"https://codeload.github.com/google/torchax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google%2Ftorchax/sbom","scorecard":{"id":1238457,"data":{"date":"2025-10-04T00:07:52Z","repo":{"name":"github.com/google/torchax","commit":"2ed1a5382d29f607eb5552a94a69b2a951ba002c"},"scorecard":{"version":"v5.1.1","commit":"cd152cb6742c5b8f2f3d2b5193b41d9c50905198"},"score":4,"checks":[{"name":"Dangerous-Workflow","score":10,"reason":"no dangerous workflow patterns detected","details":null,"documentation":{"short":"Determines if the project's GitHub Action workflows avoid dangerous patterns.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#dangerous-workflow"}},{"name":"Code-Review","score":0,"reason":"Found 0/5 approved changesets -- score normalized to 0","details":null,"documentation":{"short":"Determines if the project requires human code review before pull requests (aka merge requests) are merged.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#code-review"}},{"name":"Binary-Artifacts","score":10,"reason":"no binaries found in the repo","details":null,"documentation":{"short":"Determines if the project has generated executable (binary) artifacts in the source repository.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#binary-artifacts"}},{"name":"Maintained","score":0,"reason":"project was created in last 90 days. please review its contents carefully","details":["Warn: Repository was created in last 90 days."],"documentation":{"short":"Determines if the project is \"actively maintained\".","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#maintained"}},{"name":"Token-Permissions","score":0,"reason":"detected GitHub workflow tokens with excessive permissions","details":["Info: topLevel permissions set to 'read-all': .github/workflows/scorecard.yml:18","Warn: no topLevel permission defined: .github/workflows/unit-tests.yml:1","Info: no jobLevel write permissions found"],"documentation":{"short":"Determines if the project's workflows follow the principle of least privilege.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#token-permissions"}},{"name":"Packaging","score":-1,"reason":"packaging workflow not detected","details":["Warn: no GitHub/GitLab publishing workflow detected."],"documentation":{"short":"Determines if the project is published as a package that others can easily download, install, easily update, and uninstall.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#packaging"}},{"name":"Vulnerabilities","score":9,"reason":"1 existing vulnerabilities detected","details":["Warn: Project is vulnerable to: PYSEC-2017-74"],"documentation":{"short":"Determines if the project has open, known unfixed vulnerabilities.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#vulnerabilities"}},{"name":"SAST","score":0,"reason":"SAST tool is not run on all commits -- score normalized to 0","details":["Warn: 0 commits out of 1 are checked with a SAST tool"],"documentation":{"short":"Determines if the project uses static code analysis.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#sast"}},{"name":"CII-Best-Practices","score":0,"reason":"no effort to earn an OpenSSF best practices badge detected","details":null,"documentation":{"short":"Determines if the project has an OpenSSF (formerly CII) Best Practices Badge.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#cii-best-practices"}},{"name":"Pinned-Dependencies","score":0,"reason":"dependency not pinned by hash detected -- score normalized to 0","details":["Warn: GitHub-owned GitHubAction not pinned by hash: .github/workflows/scorecard.yml:76: update your workflow using https://app.stepsecurity.io/secureworkflow/google/torchax/scorecard.yml/main?enable=pin","Warn: GitHub-owned GitHubAction not pinned by hash: .github/workflows/unit-tests.yml:22: update your workflow using https://app.stepsecurity.io/secureworkflow/google/torchax/unit-tests.yml/main?enable=pin","Warn: GitHub-owned GitHubAction not pinned by hash: .github/workflows/unit-tests.yml:25: update your workflow using https://app.stepsecurity.io/secureworkflow/google/torchax/unit-tests.yml/main?enable=pin","Warn: containerImage not pinned by hash: examples/train_llama_torchtitan/Dockerfile:17: pin your Docker image by updating python:3.10-slim-bullseye to python:3.10-slim-bullseye@sha256:f1fb49e4d5501ac93d0ca519fb7ee6250842245aba8612926a46a0832a1ed089","Warn: pipCommand not pinned by hash: examples/train_llama_torchtitan/Dockerfile:32","Warn: pipCommand not pinned by hash: examples/train_llama_torchtitan/Dockerfile:33","Warn: pipCommand not pinned by hash: examples/train_llama_torchtitan/Dockerfile:34","Warn: pipCommand not pinned by hash: examples/train_llama_torchtitan/Dockerfile:39","Warn: pipCommand not pinned by hash: examples/train_llama_torchtitan/Dockerfile:40","Warn: pipCommand not pinned by hash: examples/train_llama_torchtitan/Dockerfile:45","Warn: pipCommand not pinned by hash: .github/workflows/unit-tests.yml:30","Warn: pipCommand not pinned by hash: .github/workflows/unit-tests.yml:31","Warn: pipCommand not pinned by hash: .github/workflows/unit-tests.yml:32","Warn: pipCommand not pinned by hash: .github/workflows/unit-tests.yml:33","Info:   2 out of   5 GitHub-owned GitHubAction dependencies pinned","Info:   1 out of   1 third-party GitHubAction dependencies pinned","Info:   0 out of   1 containerImage dependencies pinned","Info:   0 out of  10 pipCommand dependencies pinned"],"documentation":{"short":"Determines if the project has declared and pinned the dependencies of its build process.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#pinned-dependencies"}},{"name":"Fuzzing","score":0,"reason":"project is not fuzzed","details":["Warn: no fuzzer integrations found"],"documentation":{"short":"Determines if the project uses fuzzing.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#fuzzing"}},{"name":"Signed-Releases","score":-1,"reason":"no releases found","details":null,"documentation":{"short":"Determines if the project cryptographically signs release artifacts.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#signed-releases"}},{"name":"Dependency-Update-Tool","score":0,"reason":"no update tool detected","details":["Warn: no dependency update tool configurations found"],"documentation":{"short":"Determines if the project uses a dependency update tool.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#dependency-update-tool"}},{"name":"CI-Tests","score":10,"reason":"1 out of 1 merged PRs checked by a CI test -- score normalized to 10","details":null,"documentation":{"short":"Determines if the project runs tests before pull requests are merged.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#ci-tests"}},{"name":"Branch-Protection","score":-1,"reason":"internal error: error during branchesHandler.setup: internal error: githubv4.Query: Resource not accessible by integration","details":null,"documentation":{"short":"Determines if the default and release branches are protected with GitHub's branch protection settings.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#branch-protection"}},{"name":"Contributors","score":0,"reason":"project has 0 contributing companies or organizations -- score normalized to 0","details":null,"documentation":{"short":"Determines if the project has a set of contributors from multiple organizations (e.g., companies).","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#contributors"}},{"name":"License","score":10,"reason":"license file detected","details":["Info: project has a license file: LICENSE:0","Info: FSF or OSI recognized license: Apache License 2.0: LICENSE:0"],"documentation":{"short":"Determines if the project has defined a license.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#license"}},{"name":"Security-Policy","score":10,"reason":"security policy file detected","details":["Info: security policy file detected: github.com/google/.github/SECURITY.md:1","Info: Found linked content: github.com/google/.github/SECURITY.md:1","Info: Found disclosure, vulnerability, and/or timelines in security policy: github.com/google/.github/SECURITY.md:1","Info: Found text in security policy: github.com/google/.github/SECURITY.md:1"],"documentation":{"short":"Determines if the project has published a security policy.","url":"https://github.com/ossf/scorecard/blob/cd152cb6742c5b8f2f3d2b5193b41d9c50905198/docs/checks.md#security-policy"}}]},"last_synced_at":"2025-10-04T00:22:22.505Z","repository_id":317927602,"created_at":"2025-10-04T00:22:22.505Z","updated_at":"2025-10-04T00:22:22.505Z"},"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":35827727,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-07-20T02:08:10.276Z","status":"online","status_checked_at":"2026-07-24T02:00:07.870Z","response_time":62,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["ai","interoperability","jax","machine-learning","pytorch","tpu"],"created_at":"2026-01-26T13:00:33.113Z","updated_at":"2026-07-24T05:01:01.898Z","avatar_url":"https://github.com/google.png","language":"Python","funding_links":[],"categories":["Libraries"],"sub_categories":["New Libraries"],"readme":"# torchax: Running PyTorch on TPU via JAX\n\nDocs page: https://google.github.io/torchax/\nDiscord Discussion Channel: https://discord.gg/JqeJqGPyzC\n\n\n![](docs/docs/assets/logo.jpeg)\n\n**torchax** is a backend for PyTorch that allows users to run\nPyTorch programs on Google Cloud TPUs. It also provides graph-level\ninteroperability between PyTorch and JAX.\n\nWith **torchax**, you can:\n* Run PyTorch code on TPUs with minimal code changes.\n* Call JAX functions from PyTorch, passing in `jax.Array`s.\n* Call PyTorch functions from JAX, passing in `torch.Tensor`s.\n* Use JAX features like `jax.grad`, `optax`, and `GSPMD` to train PyTorch\n  models.\n* Use a PyTorch model as a feature extractor with a JAX model.\n\n## Install\n\nFirst, install the CPU version of PyTorch:\n\n```bash\n# On Linux\npip install torch --index-url https://download.pytorch.org/whl/cpu\n\n# On Mac\npip install torch\n```\n\nNext, install JAX for your desired accelerator:\n\n```bash\n# On Google Cloud TPU\npip install -U jax[tpu]\n\n# On GPU machines\npip install -U jax[cuda12]\n\n# On Linux CPU machines or Macs (see the note below)\npip install -U jax\n```\n\nNote: For Apple devices, you can install the [Metal version](https://developer.apple.com/metal/jax/) of JAX for\nhardware acceleration.\n\nFinally, install torchax:\n\n```bash\n# Install from PyPI\npip install torchax\n\n# Or, install torchax from source.\npip install git+https://github.com/google/torchax\n```\n\n## Running a Model\n\nTo execute a model with torchax, start with any `torch.nn.Module`.\nHere’s an example with a simple 2-layer model:\n\n```python\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass MyModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.fc1 = nn.Linear(28 * 28, 120)\n        self.fc2 = nn.Linear(120, 84)\n        self.fc3 = nn.Linear(84, 10)\n\n    def forward(self, x):\n        x = x.view(-1, 28 * 28)\n        x = F.relu(self.fc1(x))\n        x = F.relu(self.fc2(x))\n        x = self.fc3(x)\n        return x\n\nm = MyModel()\n\n# Execute this model using torch.\ninputs = torch.randn(3, 3, 28, 28)\nprint(m(inputs))\n```\n\nTo execute this model with `torchax`, we need to enable torchax to capture PyTorch ops:\n\n```python\nimport torchax\ntorchax.enable_globally()\n```\n\nThen, we can use a `jax` device:\n\n```python\ninputs = torch.randn(3, 3, 28, 28, device='jax')\nm = MyModel().to('jax')\nres = m(inputs)\nprint(type(res))  # outputs torchax.tensor.Tensor\nprint(res.jax()) # print the underlying Jax Array\n```\n\n`torchax.tensor.Tensor` is a `torch.Tensor` subclass that holds\na `jax.Array`. You can inspect that JAX array with `res.jax()`.\n\nAlthough the code appears to be standard PyTorch, it's actually running on JAX.\n\n## How It Works\n\ntorchax uses a `torch.Tensor` subclass, `torchax.tensor.Tensor`, which holds a\n`jax.Array` and overrides the `__torch_dispatch__` method. When a PyTorch operation\nis executed within the torchax environment (enabled by `torchax.enable_globally()`),\nthe implementation of that operation is swapped with its JAX equivalent.\n\nWhen a model is instantiated, tensor constructors like `torch.rand` create\n`torchax.tensor.Tensor` objects containing `jax.Arrays`. Subsequent operations\nextract the `jax.Array`, call the corresponding JAX implementation, and wrap the\nresult back into a `torchax.tensor.Tensor`.\n\nFor more details, see the [How It Works](docs/docs/user_guide/how-it-works.md) and\n[Ops Registry](docs/ops_registry.md) documentation.\n\n### Executing with `jax.jit`\n\nWhile torchax can run models in eager mode, `jax.jit` can be used for better performance.\n`jax.jit` is a decorator that compiles a function that takes and returns `torch.Tensors`\ninto a faster, JAX-compiled version.\n\nTo use `jax.jit`, you first need a functional version of your model where parameters\nare passed as inputs:\n\n```python\ndef model_func(param, inputs):\n  return torch.func.functional_call(m, param, inputs)\n```\n\nHere we use [torch.func.functional_call](https://pytorch.org/docs/stable/generated/torch.func.functional_call.html)\nfrom PyTorch to replace the model weights with `param` and then call the\nmodel. This is roughly equivalent to:\n\n```python\ndef model_func(param, inputs):\n  m.load_state_dict(param)\n  return m(*inputs)\n```\n\nNow, we can apply `jax_jit` on `module_func`:\n\n```python\nfrom torchax.interop import jax_jit\n\nmodel_func_jitted = jax_jit(model_func)\nprint(model_func_jitted(new_state_dict, inputs))\n```\n\nSee more examples at [eager_mode.py](examples/eager_mode.py) and the\n[examples folder](examples/).\n\nTo ease the idiom of creating functional model and calling it with parameters,\nwe also created the `JittableModule` helper class. It lets us rewrite the\nabove as:\n\n```python\nfrom torchax.interop import JittableModule\n\nm_jitted = JittableModule(m)\nres = m_jitted(...)\n```\n\nThe first time `m_jitted` is called, it will trigger `jax.jit` to compile the\ncompile for the given input shapes. Subsequent calls with the same input shapes\nwill be fast as the compilation is cached.\n\n## Saving and Loading Checkpoints\n\nYou can save and load your training state using `torchax.save_checkpoint` and `torchax.load_checkpoint`.\nThe state can be a dictionary containing the model's weights, optimizer state, and any other relevant\ninformation.\n\n```python\nimport torchax\nimport torch\nimport optax\n\n# Assume model, optimizer, and other states are defined\nmodel = MyModel()\noptimizer = optax.adam(1e-3)\nopt_state = optimizer.init(model.parameters())\nweights = model.parameters()\nbuffers = model.buffers()\nepoch = 10\n\nstate = {\n    'weights': weights,\n    'buffers': buffers,\n    'opt_state': opt_state,\n    'epoch': epoch,\n}\n\n# Save checkpoint\ntorchax.save_checkpoint(state, '/path/to/checkpoint.pt')\n\n# Load checkpoint\nloaded_state = torchax.load_checkpoint('/path/to/checkpoint.pt')\n\n# Restore state\nmodel.load_state_dict(loaded_state['weights'])\nopt_state = loaded_state['opt_state']\nepoch = loaded_state['epoch']\n```\n\n## Citation\n\n```\n@software{torchax,\n  author = {Han Qi, Chun-nien Chan, Will Cromar, Manfei Bai, Kevin Gleanson},\n  title = {torchax: PyTorch on TPU and JAX interoperability},\n  url = {https://github.com/pytorch/xla/tree/master/torchax}\n  version = {0.0.4},\n  date = {2025-02-24},\n}\n```\n\n## Maintainers \u0026 Contributors\n\nThis library is maintained by a team within Google Cloud. It has benefited from\nmany contributions from both inside and outside the team.\n\nThank you to recent contributors.\n\n```\nHan Qi (qihqi), PyTorch/XLA\nManfei Bai (manfeibai), PyTorch/XLA\nWill Cromar (will-cromar), Meta\nMilad Mohammadi (miladm), PyTorch/XLA\nSiyuan Liu (lsy323), PyTorch/XLA\nBhavya Bahl (bhavya01), PyTorch/XLA\nPei Zhang (zpcore), PyTorch/XLA\nYifei Teng (tengyifei), PyTorch/XLA\nChunnien Chan (chunnienc), Google, ODML\nAlban Desmaison (albanD), Meta, PyTorch\nSimon Teo (simonteozw), Google (20%)\nDavid Huang (dvhg), Google (20%)\nBarni Seetharaman (barney-s), Google (20%)\nAnish Karthik (anishfish2), Google (20%)\nYao Gu (guyao), Google (20%)\nYenkai Wang (yenkwang), Google (20%)\nGreg Shikhman (commander), Google (20%)\nMatin Akhlaghinia (matinehAkhlaghinia), Google (20%)\nTracy Chen (tracych477), Google (20%)\nMatthias Guenther (mrguenther), Google (20%)\nWenXin Dong (wenxindongwork), Google (20%)\nKevin Gleason (GleasonK), Google, StableHLO\nNupur Baghel (nupurbaghel), Google (20%)\nGwen Mittertreiner (gmittert), Google (20%)\nZeev Melumian (zmelumian), Lightricks\nVyom Sharma (vyom1611), Google (20%)\nShitong Wang (ShitongWang), Adobe\nRémi Doreau (ayshiff), Google (20%)\nLance Wang (wang2yn84), Google, CoreML\nHossein Sarshar (hosseinsarshar), Google (20%)\nDaniel Vega-Myhre (danielvegamyhre), Google (20%)\nTianqi Fan (tqfan28), Google (20%)\nJim Lin (jimlinntu), Google (20%)\nFanhai Lu (FanhaiLu1), Google Cloud\nDeWitt Clinton (dewitt), Google PyTorch\nAman Gupta (aman2930), Google (20%)\n```\n\nA special thank you to @albanD for the [initial inspiration](https://github.com/albanD/subclass_zoo/blob/main/new_device.py)\nfor torchax.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle%2Ftorchax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fgoogle%2Ftorchax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle%2Ftorchax/lists"}