{"categories":[{"label":"Software Development","url":"https://skillfed.io/packages/category/software-development/14"},{"label":"Libraries","url":"https://skillfed.io/packages/category/software-development-libraries/8"},{"label":"Scientific/Engineering","url":"https://skillfed.io/packages/category/scientific-engineering/6"},{"label":"Python Modules","url":"https://skillfed.io/packages/category/software-development-libraries-python-modules/17"},{"label":"Artificial Intelligence","url":"https://skillfed.io/packages/category/scientific-engineering-artificial-intelligence/8"},{"label":"Mathematics","url":"https://skillfed.io/packages/category/scientific-engineering-mathematics/3"}],"enrichment":{"capability":"torchax is a PyTorch backend that runs PyTorch code on Google Cloud TPUs via JAX, with bidirectional interoperability allowing you to call JAX functions from PyTorch and vice versa.","skillfed_tags":["tpu-backend","pytorch-jax-interop","accelerator-support"],"use_cases":["Run existing PyTorch models on Google Cloud TPUs without rewriting them for JAX.","Use JAX optimization libraries and gradient transformations to train PyTorch models.","Call JAX functions with compiled performance from within PyTorch code, passing jax.Arrays directly.","Combine a PyTorch feature extractor with a JAX model in the same training loop.","Leverage JAX features like GSPMD for distributed training of PyTorch models."],"what_it_does":"torchax bridges PyTorch and JAX by implementing a custom PyTorch backend that intercepts PyTorch operations and executes them via JAX. This allows you to write standard PyTorch code but run it on TPUs and other JAX-supported accelerators, and to mix PyTorch and JAX code in the same program\u2014calling JAX functions with tensors and vice versa.\n\nThe package works by subclassing torch.Tensor into torchax.tensor.Tensor and overriding __torch_dispatch__ to redirect operations to JAX equivalents. You enable it globally with torchax.enable_globally(), then create tensors and models on the 'jax' device. Operations run in eager mode by default, but you can wrap functions with jax.jit for compiled execution. The package also provides checkpoint save/load utilities and a JittableModule helper for easier JIT compilation of models.","worth_installing":"Yes, if you need to run PyTorch on TPUs or want deep PyTorch\u2013JAX interoperability. The package is actively maintained, has no known vulnerabilities, and installs easily. However, it is in alpha (Development Status :: 3), so expect incomplete operation coverage and potential API changes. Verify that your specific operations are supported before committing to production use."},"id":"torchax","links":{"html":"https://skillfed.io/packages/torchax","md":"https://skillfed.io/packages/torchax.md","pypi":"https://pypi.org/project/torchax/"},"maintenance":{"status":"active"},"meta":{"latest_release":"2026-06-17","license_spdx":null,"license_treatment":"permissive","name":"torchax","python_support":"supports_current","summary":"torchax is a library for running Jax and PyTorch together"},"popularity":{"monthly_downloads":164579,"position":10542,"tier":"top_15000"},"security":{"n_vulnerabilities":0},"version":"0.0.13"}