torchax 0.0.4__py3-none-any.whl

torchax/train.py

@@ -0,0 +1,120 @@
+import collections
+import functools
+import torch
+import jax
+import torchax
+from torchax import interop
+from torchax.interop import torch_view, jax_view
+import optax
+
+
+remat = torch_view(jax.remat)
+mark_sharding = torch_view(jax.lax.with_sharding_constraint)
+
+
+def make_train_step(model_fn, 
+                    loss_fn, optax_optimizer, 
+                    remat_policy=None):
+  """Make a function that do one train step given model and loss.
+
+  model_fn: a function representing the model's forward:
+      i.e. has signature Callable[weights, buffers, args] -> result. Where,
+      weights is a pytree of trainable parameters
+      buffers is a pytree of non-trainable parameters / constants
+      args is the input data loaded from the data set
+      result is the return value of the model
+  loss_fn: a function to compute loss.
+      i.e. it has signature of Callable[result, label] -> loss
+      where, result is what model_fn returned
+        loss is loaded from the dataloader.
+  optax_optimizer: the optimizer from optax library. for example, optax.adam
+  remat_policy: One of jax.ad_checkpoint.checkpoint_policies, specifies how
+      to do gradient checkpointing. If None, then it means checkpoint everything.
+  """
+  env = torchax.default_env()
+  def loss(weights, buffers, args, label): # inputs are XLATensor
+    with env, jax.named_scope('compute_loss'):
+      res = model_fn(weights, buffers, args)
+      l = loss_fn(res, label)
+      return l
+
+  loss = interop.gradient_checkpoint(loss, kwargs={'policy': remat_policy})
+  grad_fn = interop.jax_value_and_grad(loss)
+
+  def step(weights, buffers, opt_state, args, label): #inputs are array
+    with jax.named_scope('compute_gradient'):
+        loss, gradient = grad_fn(weights, buffers, args, label)
+
+    with jax.named_scope("optimizer_updates"):
+        updates, opt_state = interop.call_jax(
+            optax_optimizer.update,
+            gradient, opt_state, weights)
+        weights = interop.call_jax(optax.apply_updates, weights, updates)
+    return loss, weights, opt_state
+
+  # TODO: apply jax.jit so the user don't have to.
+  return step
+
+  
+  
+  
+class Container:
+  pass
+
+class ScannedModule(torch.nn.Module):
+
+  def __init__(self, module_list, checkpoint_policy=None):
+    super().__init__()
+
+    self.c = None
+    assert module_list
+    self.c = Container()
+    self.c.one_mod = module_list[0]
+    self.checkpoint_policy = checkpoint_policy
+
+    weights = self._stack_layer_weights(module_list)
+    self.layer_weights_keys = list(self.c.one_mod.state_dict().keys())
+    self.params = torch.nn.ParameterDict({
+      self._param_name_new(k): v for k, v in weights.items()
+    })
+  
+  def _stack_layer_weights(self, module_list):
+    # Create weights such that, for every [n, m] weights
+    # becomes [k, n, m] where k is number of layer
+    # i.e. stacking layer weights together
+    temp = collections.defaultdict(list)
+    for m in module_list:
+      for k, v in m.state_dict().items():
+          temp[k].append(v)
+    res = {k: torch.stack(v) for k, v in temp.items()}
+    return res
+
+
+  def _param_name_new(self, old):
+      return '___'.join(old.split('.'))
+
+  def _param_name_old(self, new):
+      return '.'.join(new.split('___'))
+
+  def forward(self, *args, **kwargs):
+      assert not kwargs
+      weights = {k: self.params[self._param_name_new(k)] for k in self.layer_weights_keys}
+      scan = interop.torch_view(jax.lax.scan)
+
+      def eval_one_layer(args, weight):
+          # unpack args
+          h, *rest = args
+          newh = torch.func.functional_call(self.c.one_mod, weight, args)
+          # next layer's input; and residual to be added to list
+          return (newh, *rest), None 
+
+      _eval_one_layer = interop.gradient_checkpoint(
+          eval_one_layer,
+          kwargs={'policy': self.checkpoint_policy},
+      )
+      h, _ = scan(
+          _eval_one_layer,
+          args,
+          weights,
+      )
+      return h[0]

torchax/types.py

@@ -0,0 +1,12 @@
+from typing import Callable, Any, Union, ParamSpec, TypeAlias
+import torch
+import jax
+import jax.numpy as jnp
+import sys
+
+P = ParamSpec('P')
+
+TorchValue: TypeAlias = Union[torch.Tensor, torch.dtype, 'TorchCallable', Any]
+TorchCallable: TypeAlias = Callable[P, TorchValue]
+JaxValue: TypeAlias = Union[jax.Array, jnp.dtype, 'JaxCallable', Any]
+JaxCallable: TypeAlias = Callable[P, JaxValue]

torchax-0.0.4.dist-info/METADATA

@@ -0,0 +1,341 @@
+Metadata-Version: 2.4
+Name: torchax
+Version: 0.0.4
+Summary: torchax is a library for running PyTorch on TPU
+Project-URL: Homepage, https://github.com/pytorch/xla/tree/master/torchax
+Author-email: Han Qi <qihan.dev@gmail.com>, Pytorch/XLA team <pytorchxla-dev@google.com>
+License: BSD 3-Clause License
+        
+        Copyright (c) 2023, pytorch-tpu
+        
+        Redistribution and use in source and binary forms, with or without
+        modification, are permitted provided that the following conditions are met:
+        
+        1. Redistributions of source code must retain the above copyright notice, this
+           list of conditions and the following disclaimer.
+        
+        2. Redistributions in binary form must reproduce the above copyright notice,
+           this list of conditions and the following disclaimer in the documentation
+           and/or other materials provided with the distribution.
+        
+        3. Neither the name of the copyright holder nor the names of its
+           contributors may be used to endorse or promote products derived from
+           this software without specific prior written permission.
+        
+        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+        AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+        IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+        DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+        FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+        DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+        SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+        CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+        OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+        OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.10
+Provides-Extra: cpu
+Requires-Dist: jax[cpu]; extra == 'cpu'
+Requires-Dist: jax[cpu]>=0.4.30; extra == 'cpu'
+Provides-Extra: cuda
+Requires-Dist: jax[cpu]>=0.4.30; extra == 'cuda'
+Requires-Dist: jax[cuda12]; extra == 'cuda'
+Provides-Extra: odml
+Requires-Dist: jax[cpu]; extra == 'odml'
+Requires-Dist: jax[cpu]>=0.4.30; extra == 'odml'
+Provides-Extra: tpu
+Requires-Dist: jax[cpu]>=0.4.30; extra == 'tpu'
+Requires-Dist: jax[tpu]; extra == 'tpu'
+Description-Content-Type: text/markdown
+
+# torchax: Running PyTorch on TPU
+
+**torchax!** is a backend for PyTorch, allowing users to run
+PyTorch on Google CloudTPUs. **torchax!** is also a library for providing
+graph-level interoperability between PyTorch and Jax.
+
+This means, with **torchax** you can:
+* Run PyTorch code on TPU with as little as 2 lines of code change.
+* Call a jax function from a pytorch function, passing in `jax.Array`s
+* Call a pytorch function from a jax function, passing in a `torch.Tensor` subclass.
+* Use jax features such as `jax.grad`, `optax` and `GSMPD` to train a Pytorch model.
+* Use a Pytorch model as feature extractor and use it with a Jax model.
+etc etc.
+
+## Install
+
+
+### On Google Cloud TPU:
+First install torch CPU:
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+```
+
+Then install jax TPU:
+
+```bash
+pip install -U jax[tpu]
+```
+
+Finally install torchax
+
+```bash
+pip install torchax
+```
+
+### On GPU machines:
+First install torch CPU:
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+```
+
+Then install jax CUDA:
+
+```bash
+pip install -U jax[cuda12]
+```
+
+Finally install torchax
+
+```bash
+pip install torchax
+```
+
+### On CPU machines (mac included)
+First install torch CPU:
+
+```bash
+# Linux
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+# OR Mac:
+pip install torch
+```
+
+Then install jax CPU:
+
+```bash
+pip install -U jax
+```
+
+Finally install torchax
+
+```bash
+pip install torchax
+```
+
+NOTE: if you like metal support for Apple devices then install the
+metal version of jax: https://developer.apple.com/metal/jax/
+
+### Installing `torchax` from source
+
+Still need to install `torch` CPU and `Jax` of your accelerator (GPU, TPU or None).
+
+```bash
+pip install git+https://github.com/pytorch/xla.git#subdirectory=torchax
+```
+
+## Run a model
+
+Now let's execute a model under torchax. We'll start with a simple 2-layer model
+it can be in theory any instance of `torch.nn.Module`.
+
+```python
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(28 * 28, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x):
+        x = x.view(-1, 28 * 28)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+m = MyModel()
+
+# Execute this model using torch
+inputs = torch.randn(3, 3, 28, 28)
+print(m(inputs))
+```
+
+This model `m` contains 2 parts: the weights that is stored inside of the model
+and it's submodules (`nn.Linear`).
+
+To execute this model with `torchax`; we need to enable torchax to capture pytorch ops.
+To enable this, use:
+
+```python
+import torchax
+torchax.enable_globally()
+```
+Then, a `jax` device will be available to use
+
+```python
+inputs = torch.randn(3, 3, 28, 28, device='jax')
+m = MyModel()
+res = m(inputs)
+print(type(res))  # outputs torchax.tensor.Tensor
+```
+
+`torchax.tensor.Tensor` is a `torch.Tensor` subclass that holds
+a `jax.Array`. You can inspect that jax array with `res.jax()`
+
+
+## What is happening behind the scene:
+
+We took the approach detailed in [new device](https://github.com/albanD/subclass_zoo/blob/main/new_device.py) recipe by Alban (@albanD); using `jax.Array` for the `raw_data`.
+
+In other words, When a torch op is executed inside of `env` context manager (which is enabled with `torchax.enable_globally()`), we can swap out the
+implementation of that op written in Jax.
+
+When a model's constructor runs, it will call some tensor constructor, such as
+`torch.rand`, `torch.ones` or `torch.zeros` etc to create its weights. The constructor
+will create an `torch.Tensor` subclass that contains a `jax.Array`.
+
+Then, each subsequent op can unpack the `jax.Array`, call the op implementation,
+and wraps it back into `torch.Tensor` subclass.
+
+See more at [how_it_works](docs/how_it_works.md) and [ops registry](docs/ops_registry.md).
+
+
+### Executing with jax.jit
+
+The above script will execute the model using eager mode Jax as backend. This
+does allow executing torch models on TPU, but is often slower than what we can
+achieve with `jax.jit`.
+
+`jax.jit` is a function that takes a Jax function (i.e. a function that takes jax array
+and returns jax array) into the same function, but faster.
+
+We have made the `jax_jit` decorator that would accomplish the same with functions
+that takes and returns `torch.Tensor`. To use this, the first step is to create
+a functional version of this model: this means the parameters should be passed in
+as input instead of being attributes on class:
+
+
+```python
+
+def model_func(param, inputs):
+  return torch.func.functional_call(m, param, inputs)
+
+```
+Here we use [torch.func.functional_call](https://pytorch.org/docs/stable/generated/torch.func.functional_call.html)
+from PyTorch to replace the model
+weights with `param`, then call the model. This is roughly equivalent to:
+
+```python
+def model_func(param, inputs):
+  m.load_state_dict(param)
+  return m(*inputs)
+```
+
+Now, we can apply `jax_jit`
+
+```python
+from torchax.interop import jax_jit
+model_func_jitted = jax_jit(model_func)
+print(model_func_jitted(new_state_dict, inputs))
+```
+
+See more examples at [eager_mode.py](examples/eager_mode.py) and the (examples folder)[examples/]
+
+However, to ease the idiom of creating functional model and calling it with parameters,
+we also created the `JittableModule` helper class.
+
+So the above can be written as:
+
+```python
+
+from torchax.interop import JittableModule
+
+m_jitted = JittableModule(m)
+res = m_jitted(...)
+```
+
+The first time that `m_jitted` is called , it will trigger `jax.jit`
+then the subsequent computation with inputs of same shape will be fast.
+
+
+
+# Citation:
+
+@software{torchax,
+  author = {Han Qi, Chun-nien Chan, Will Cromar, Manfei Bai, Kevin Gleanson},
+  title = {torchax: PyTorch on TPU and Jax interoperability},
+  url = {https://github.com/pytorch/xla/tree/master/torchax}
+  version = {0.0.4},
+  date = {2025-02-24},
+}
+
+# Maintainers & Contributors:
+
+This library is created and maintained by the PyTorch/XLA team at Google Cloud.
+
+However, it benefitted from many direct and indirect
+contributions outside of the team. Many of them done by
+fellow Googlers using [Google's 20% project policy](https://ebsedu.org/blog/google-tapping-workplace-actualization-20-time-rule), others by partner teams.
+
+Here is the full list of contributors by 2025-02-25.
+
+Han Qi (qihqi), Pytorch / XLA
+Manfei Bai (manfeibai), Pytorch / XLA
+Will Cromar (will-cromar), Meta
+Milad Mohammadi (miladm), Pytorch / XLA
+Siyuan Liu (lsy323), Pytorch / XLA
+Bhavya Bahl (bhavya01), Pytorch / XLA
+Pei Zhang (zpcore), Pytorch / XLA
+Yifei Teng (tengyifei), Pytorch / XLA
+Chunnien Chan (chunnienc), Google, ODML
+Alban Desmaison (albanD), Meta, Pytorch
+Simon Teo (simonteozw), Google(20%)
+David Huang (dvhg), Google(20%)
+Barni Seetharaman (barney-s), Google(20%)
+Anish Karthik (anishfish2) , Google(20%)
+Yao Gu (guyao) , Google(20%)
+Yenkai Wang (yenkwang) , Google(20%)
+Greg Shikhman (commander) , Google(20%)
+Matin Akhlaghinia (matinehAkhlaghinia), Google(20%)
+Tracy Chen (tracych477), Google(20%)
+Matthias Guenther (mrguenther) , Google(20%)
+WenXin Dong (wenxindongwork), Google(20%)
+Kevin Gleason (GleasonK) , Google, StableHLO
+Nupur Baghel (nupurbaghel), Google(20%)
+Gwen Mittertreiner (gmittert), Google(20%)
+Zeev Melumian (zmelumian), Lightricks
+Vyom Sharma (vyom1611), Google(20%)
+Shitong Wang (ShitongWang), Adobe
+Rémi Doreau (ayshiff), Google(20%)
+Lance Wang (wang2yn84), Google, CoreML
+Hossein Sarshar (hosseinsarshar) , Google(20%)
+Daniel Vega-Myhre (danielvegamyhre) , Google(20%)
+Tianqi Fan (tqfan28), Google(20%)
+Jim Lin (jimlinntu), Google(20%)
+Fanhai Lu (FanhaiLu1), Google Cloud
+DeWitt Clinton (dewitt), Google PyTorch
+Aman Gupta (aman2930) , Google(20%)

torchax-0.0.4.dist-info/RECORD

@@ -0,0 +1,27 @@
+torchax/CONTRIBUTING.md,sha256=T6Uy7vZvzXArbPqqUswEtQBgTU_eRwtJnHOiZNth2Pw,1371
+torchax/__init__.py,sha256=0NoxpzaPIMYMJXZz_JiRCTYZCuICnTFrLDm5cZKGIm0,3242
+torchax/config.py,sha256=Zpfcn7Q3QsKnn4cCx1-bYxA8q_lqrWJcUz30B3QcIc4,535
+torchax/decompositions.py,sha256=C067xSd9MpGs1SDqlC5T9RnqNj-i-m31mysfFGDxLEQ,11221
+torchax/device_module.py,sha256=NNEGSfk9ApVVSy5dRjHkvGuNGamopCEjSlItcNlAIbE,253
+torchax/distributed.py,sha256=CXkpV0K2Oas7fcRA-i3s2VXTz8l4CEI7_NnbYOsHtRw,7519
+torchax/environment.py,sha256=daEdpEyAJIa8b2VkCqSKcw8PaExcB6Qro80XNes_sHA,2
+torchax/export.py,sha256=MG19Y0QcRcX3S3gruy9K-wFGdemD4I0ii03HWQM31xk,8853
+torchax/interop.py,sha256=fncREqC1Nib-ud7wyekmtKaif7GSDql7ITmzTt5A9e4,6774
+torchax/tensor.py,sha256=89yXmyO1uT6F1KI-Kn6NCnIieCiiMsDutgS071_9p64,16644
+torchax/tf_integration.py,sha256=d_h4vSJm7N9rJXpUPNCDOiUz3J1-UPo3KU8D9Wi4nnc,4074
+torchax/train.py,sha256=Ym_SC0WOmhQMhSCH0aZ-zgP7Q19WdCDTeoI-JBFrqx8,3930
+torchax/types.py,sha256=j4ERjkgDgwhgi9zrwwbbiv4HMDlrJ1IEMUCmP_BIJ9M,388
+torchax/ops/__init__.py,sha256=Vr1p8zDHwfXZBUbw70iNiCJLZLNdI6gR_vUlaiA7Usg,270
+torchax/ops/jaten.py,sha256=ycTSyCzgJpz5zu936BTZXmRJWuzhxeGOOjoffU00c3g,154999
+torchax/ops/jax_reimplement.py,sha256=29SGPF0bszgCuyVxiZ4A2xTPtAuPXBDvDG79M3C0Vxo,7468
+torchax/ops/jc10d.py,sha256=DiJjiUFHyYmFUHEgqEDbGsLaCsJs94xJlBatRwXiPfg,1317
+torchax/ops/jlibrary.py,sha256=R22GUlK6lUivuDSmvtqwbV9ku_uT0ddLcQcwHI96vqM,2927
+torchax/ops/jtorch.py,sha256=fWU2z5uD3FpVE5ui-O5nfHEule2hod4BqCnsbV5jVfQ,14463
+torchax/ops/jtorchvision_nms.py,sha256=SiMN9bUV7GADD2S2PkH9spn5oQGgr2jhFQ1nyd3eHgI,8873
+torchax/ops/mappings.py,sha256=a8C5CYIEucGIxUPQHCsDZypBC3b_BYtOKQTWf9bdiRc,3060
+torchax/ops/op_base.py,sha256=EdlAx3oqU2brMngjZvTxECin0eNiPdAUfOLRtFL3tgg,2900
+torchax/ops/ops_registry.py,sha256=nXCAXiQNpC0uBUl8xEim7TOUTyROZm_4H1LaOvYwJVQ,1241
+torchax-0.0.4.dist-info/METADATA,sha256=4VQb39quiN9U9j955fs9zY00x4t4Orubo92ONCCLyug,11119
+torchax-0.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+torchax-0.0.4.dist-info/licenses/LICENSE,sha256=ZHyir3-ltOerFLt9JH1bjf7lIxIWipFmqeMnB_8z_aU,1498
+torchax-0.0.4.dist-info/RECORD,,

torchax-0.0.4.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

torchax-0.0.4.dist-info/licenses/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2023, pytorch-tpu
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.