torchax 0.0.6__py3-none-any.whl → 0.0.10.dev20251116__py3-none-any.whl

torchax-0.0.10.dev20251116.dist-info/licenses/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

torchax/configuration.py

@@ -1,30 +0,0 @@
-import dataclasses
-
-
-@dataclasses.dataclass
-class Configuration:
-  debug_print_each_op: bool = False
-  debug_accuracy_for_each_op: bool = False
-  debug_mixed_tensor: bool = False
-  debug_print_each_op_operands: bool = False
-
-  use_int32_for_index: bool = False
-
-  # normally, math between CPU torch.Tensor with torchax.Tensor is not
-  # allowed. However, if that torch.Tensor happens to be scalar, then we
-  # can use scalar * tensor math to handle it
-  allow_mixed_math_with_scalar_tensor: bool = True
-
-  # If true, we will convert Views into torchax.Tensors eagerly
-  force_materialize_views: bool = False
-
-  # Use DLPack for converting jax.Arrays <-> and torch.Tensor
-  use_dlpack_for_data_conversion: bool = False
-
-  # Flash attention
-  use_tpu_flash_attention: bool = False
-  shmap_flash_attention: bool = False
-
-  # device
-  treat_cuda_as_jax_device: bool = True
-  internal_respect_torch_return_dtypes: bool = False

torchax/environment.py

@@ -1 +0,0 @@
-

torchax/tf_integration.py

@@ -1,119 +0,0 @@
-# pylint: disable
-import os
-from typing import Any, Tuple
-
-from jax.experimental import jax2tf
-import tensorflow as tf
-import torch
-from torchax import export
-
-
-def exported_program_to_tf_function(ep, enable_xla=True):
-  weights, jax_program = export.exported_program_to_jax(ep)
-  wrapped = lambda *args: jax_program(weights, (args,))
-  avals = export.extract_avals(ep)
-  input_signature = [
-      tf.TensorSpec(shape=t.shape, dtype=t.dtype, name=f"args_{i}")
-      for i, t in enumerate(avals)
-  ]
-  tf_f = tf.function(
-      jax2tf.convert(
-          wrapped,
-          with_gradient=False,
-          enable_xla=enable_xla,
-      ),
-      autograph=False,
-      input_signature=input_signature,
-  )
-  return tf_f
-
-
-def exported_program_to_tf_module(ep: torch.export.ExportedProgram,
-                                  enable_xla=True) -> tf.Module:
-  tfm = tf.Module()
-  tfm.f = exported_program_to_tf_function(ep, enable_xla)
-  return tfm
-
-
-def save_exported_program_as_tf_saved_model(
-    ep: torch.export.ExportedProgram,
-    saved_model_dir: os.PathLike,
-    serving_key: str = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
-    function_alias: str = "",
-    enable_xla=True,
-):
-  """This function will export and save a pytorch ExportedProgram to tf.saved_model format.
-
-  The resulting tf.saved_model can be used inference using tf.serving model
-  server
-  or further convert to tflite flatbuffer for on-device serving.
-
-  Args:
-    torch_model: torch.nn.Module - model to export and save
-    args: Tuple[Any] - a set of args to trace the model with, i.e.
-      torch_model(*args) must run
-    saved_model_dir: os.PathLike - location to an empty directory to store the
-      saved_model
-    serving_key: str  - serving key tag, this is used by tf.serving to know
-      which function to run.
-    function_alias: str - passed through saved_model.save, used to tag a
-      function for inference converter or other tools.
-  """
-  tfm = exported_program_to_tf_module(ep, enable_xla=enable_xla)
-  signatures = {
-      serving_key: tfm.f.get_concrete_function(*tfm.f.input_signature)
-  }
-  save_options = tf.saved_model.SaveOptions(function_aliases={
-      function_alias: tfm.f,
-  })
-  tf.saved_model.save(
-      tfm,
-      saved_model_dir,
-      signatures=signatures,
-      options=save_options,
-  )
-
-
-def save_torch_module_as_tf_saved_model(
-    torch_model: torch.nn.Module,
-    args: Tuple[Any],
-    saved_model_dir: os.PathLike,
-    serving_key: str = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
-    function_alias: str = "",
-    enable_xla=True,
-):
-  """This function will export and save a pytorch nn.Module to tf.saved_model format.
-
-  The resulting tf.saved_model can be used inference using tf.serving model
-  server
-  or further convert to tflite flatbuffer for on-device serving.
-
-  Args:
-    torch_model: torch.nn.Module - model to export and save
-    args: Tuple[Any] - a set of args to trace the model with, i.e.
-      torch_model(*args) must run
-    saved_model_dir: os.PathLike - location to an empty directory to store the
-      saved_model
-    serving_key: str  - serving key tag, this is used by tf.serving to know
-      which function to run.
-    function_alias: str - passed through saved_model.save, used to tag a
-      function for inference converter or other tools.
-  """
-  ep = torch.export.export(torch_model, args)
-  save_exported_program_as_tf_saved_model(ep, saved_model_dir, serving_key,
-                                          function_alias, enable_xla)
-
-
-def exported_program_to_tflite_flatbuffer(ep: torch.export.ExportedProgram):
-  tfm = exported_program_to_tf_module(ep)
-  tf_concrete_func = tfm.f.get_concrete_function(*tfm.f.input_signature)
-  converter = tf.lite.TFLiteConverter.from_concrete_functions(
-      [tf_concrete_func], tfm)
-  tflite_model = converter.convert()
-  return tflite_model
-
-
-def torch_module_to_tflite_flatbuffer(torch_model: torch.nn.Module,
-                                      args: Tuple[Any]):
-  ep = torch.export.export(torch_model, args)
-  return exported_program_to_tflite_flatbuffer(ep)

torchax-0.0.6.dist-info/METADATA

@@ -1,307 +0,0 @@
-Metadata-Version: 2.4
-Name: torchax
-Version: 0.0.6
-Summary: torchax is a library for running Jax and PyTorch together
-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.6.2; extra == 'cpu'
-Provides-Extra: cuda
-Requires-Dist: jax[cpu]>=0.6.2; extra == 'cuda'
-Requires-Dist: jax[cuda12]; extra == 'cuda'
-Provides-Extra: odml
-Requires-Dist: jax[cpu]; extra == 'odml'
-Requires-Dist: jax[cpu]>=0.6.2; extra == 'odml'
-Provides-Extra: tpu
-Requires-Dist: jax[cpu]>=0.6.2; extra == 'tpu'
-Requires-Dist: jax[tpu]; extra == 'tpu'
-Description-Content-Type: text/markdown
-
-# torchax: Running PyTorch on TPU via JAX
-
-**torchax** is a backend for PyTorch, allowing users to run
-PyTorch on Google Cloud TPUs. **torchax** is also a library for providing
-graph-level interoperability between PyTorch and JAX.
-
-This means, with **torchax** you can:
-* Run PyTorch code on TPUs 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`s.
-* Use JAX features such as `jax.grad`, `optax`, and `GSPMD` to train a PyTorch
-  model.
-* Use a PyTorch model as feature extractor and use it with a JAX model.
-etc etc.
-
-## Install
-
-First install torch CPU:
-
-```bash
-# On Linux.
-pip install torch --index-url https://download.pytorch.org/whl/cpu
-
-# Or on Mac.
-pip install torch
-```
-
-Then install JAX for the accelerator you want to use:
-
-```bash
-# On Google Cloud TPU.
-pip install -U jax[tpu]
-
-# Or, on GPU machines.
-pip install -U jax[cuda12]
-
-# Or, on Linux CPU machines or Macs (see the note below).
-pip install -U jax
-```
-
-NOTE: if you like metal support for Apple devices then install the
-metal version of JAX: https://developer.apple.com/metal/jax/
-
-Finally install torchax:
-
-```bash
-# Install pre-built torchax.
-pip install torchax
-
-# Or, install torchax from source.
-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.
-In theory, we can use 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))
-```
-
-To execute this model with `torchax`, we need to enable torchax to capture PyTorch ops:
-
-```python
-import torchax
-torchax.enable_globally()
-```
-
-Then, we can use a `jax` device:
-
-```python
-inputs = torch.randn(3, 3, 28, 28, device='jax')
-m = MyModel().to('jax')
-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 the
-[new device](https://github.com/albanD/subclass_zoo/blob/main/new_device.py)
-recipe by Alban (@albanD), using `jax.Array` for `raw_data`.
-
-In other words, when a torch op is executed inside an `env` context manager,
-which is enabled by `torchax.enable_globally()`, we will swap out the
-implementation of that op with JAX.
-
-When a model's constructor runs, it will call some tensor constructor, such as
-`torch.rand`, `torch.ones`, or `torch.zeros` to create its weights. When torchax
-is enabled, these constructors will create a `torchax.tensor.Tensor`, which
-contains a `jax.Array`.
-
-Then, each subsequent op will extract the `jax.Array`, call the op's JAX
-implementation, and wrap the result back into a `torchax.tensor.Tensor`,
-
-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 the backend. This
-does allow executing torch models on TPUs, 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 arrays
-and returns JAX arrays) into a compiled (thus faster) version of the same function.
-
-We have made a `jax_jit` decorator that would accomplish the same with functions
-that takes and returns `torch.Tensor`s. 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 of the 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` and 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` on `module_func`:
-
-```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/).
-
-To ease the idiom of creating functional model and calling it with parameters,
-we also created the `JittableModule` helper class. It lets us rewrite the
-above as:
-
-```python
-from torchax.interop import JittableModule
-
-m_jitted = JittableModule(m)
-res = m_jitted(...)
-```
-
-The first time `m_jitted` is called, it will trigger `jax.jit` to compile the
-compile for the given input shapes. Subsequent calls with the same input shapes
-will be fast as the compilation is cached.
-
-## 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.
-
-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 at Google and other companies.
-
-Here is the 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%)
-```