Revert 2 accidental commits that I made. (#9536)

Author: qihqi

torchax/torchax/CONTRIBUTING.md

@@ -0,0 +1,38 @@
+# Contributing to TorchXLA2
+
+We appreciate all contributions. If you are planning to contribute a bug fix for an open issue, please comment on the thread and we're happy to provide any guidance. You are very welcome to pick issues from good first issue and help wanted labels.
+
+If you plan to contribute new features, utility functions or extensions to the core, please first open an issue and discuss the feature with us. Sending a PR without discussion might end up resulting in a rejected PR, because we might be taking the core in a different direction than you might be aware of.
+
+
+# Developer setup
+
+## Mac setup:
+@qihqi
+
+I am able to develop directly on mac (m1) laptop for most of parts. Using steps
+in README.md works. The condensed version for easy copy & paste:
+
+```bash
+conda create --name <your_name> python=3.10
+conda activate <your_name>
+pip install --upgrade "jax[cpu]" torch
+pip install -r test_requirements.txt
+pip install -e .
+pytest test
+```
+
+### VSCode
+
+I use vscode on my Mac. I loosely followed instruction in
+https://code.visualstudio.com/docs/python/python-tutorial
+to setup a proper python environment.
+
+The plugins I installed (a subset of the ones listed above) are:
+* VSCode's official Python plugin
+* Ruff formatter
+* Python Debugger
+
+I also changed Python interpreter to point at the one in my conda env.
+That is all the changes I have.
+

torchax/torchax/__init__.py

@@ -40,30 +40,7 @@ def default_env():
 
 
 def extract_jax(mod: torch.nn.Module, env=None):
-  """Extracts the state of a `torch.nn.Module` into a JAX-compatible format.
-
-  **Arguments:**
-
-  *   `mod` (`torch.nn.Module`): The PyTorch model to extract the state from.
-  *   `env` (optional): The `torchax` environment to use. If not provided, the default environment is used.
-
-  **Returns:**
-
-  A tuple containing:
-
-  *   A `pytree` of `jax.ndarray` representing the model's state (parameters and buffers).
-  *   A JAX-callable function that executes the model's forward pass.
-
-  **Usage:**
-
-  ```python
-  import torch
-  import torchax
-
-  model = torch.nn.Linear(10, 20)
-  states, jax_func = torchax.extract_jax(model)
-  ```
-  """
+  """Returns a pytree of jax.ndarray and a jax callable."""
   if env is None:
     env = default_env()
   states = dict(mod.named_buffers())
@@ -83,31 +60,11 @@ def jax_func(states, args, kwargs=None):
 
 
 def enable_globally():
-  """Enables `torchax` globally, which intercepts PyTorch operations and routes them to the JAX backend. This is the primary entry point for using `torchax`.
-
-  **Usage:**
-
-  ```python
-  import torchax
-
-  torchax.enable_globally()
-  ```
-  """
   env = default_env().enable_torch_modes()
   return env
 
 
 def disable_globally():
-  """Disables the `torchax` backend. After calling this, PyTorch operations will revert to their default behavior.
-
-  **Usage:**
-
-  ```python
-  import torchax
-
-  torchax.disable_globally()
-  ```
-  """
   global env
   default_env().disable_torch_modes()
 
@@ -153,40 +110,6 @@ class CompileOptions:
 
 
 def compile(fn, options: Optional[CompileOptions] = None):
-  """Compiles a function or `torch.nn.Module` for optimized execution with JAX.
-
-  **Arguments:**
-
-  *   `fn`: The function or `torch.nn.Module` to compile.
-  *   `options` (`CompileOptions`, optional): A `CompileOptions` object to configure the compilation process.
-
-  **`CompileOptions`:**
-
-  *   `methods_to_compile` (`List[str]`, default=`['forward']`): A list of methods to compile when `fn` is a `torch.nn.Module`.
-  *   `jax_jit_kwargs` (`Dict[str, Any]`, default=`{}`): A dictionary of keyword arguments to pass to `jax.jit`.
-  *   `mode` (`str`, default=`'jax'`): The compilation mode. Currently, only `'jax'` is supported.
-
-  **Returns:**
-
-  A compiled version of the input function or module.
-
-  **Usage:**
-
-  ```python
-  import torch
-  import torchax
-
-  model = torch.nn.Linear(10, 20)
-  compiled_model = torchax.compile(model)
-
-  # With options
-  options = torchax.CompileOptions(
-      methods_to_compile=['forward', 'encode'],
-      jax_jit_kwargs={'static_argnums': (0,)}
-  )
-  compiled_model = torchax.compile(model, options)
-  ```
-  """
   options = options or CompileOptions()
   if options.mode == 'jax':
     from torchax import interop

torchax/torchax/amp.py

@@ -57,19 +57,6 @@ def is_float(a):
 
 @contextlib.contextmanager
 def autocast(device, dtype=torch.bfloat16, env=None):
-  """A context manager for automatic mixed precision (AMP).
-
-  This context manager enables automatic mixed precision, which can improve
-  performance by using lower-precision data types for certain operations.
-
-  **Arguments:**
-
-  *   `device`: The device to use for autocasting (e.g., "cuda", "cpu").
-  *   `dtype` (`torch.dtype`, optional): The lower-precision data type to use.
-      Defaults to `torch.bfloat16`.
-  *   `env` (optional): The `torchax` environment. If not provided, the default
-      environment is used.
-  """
   del device
   if env is None:
     import torchax

torchax/torchax/config.py

@@ -3,36 +3,6 @@
 
 @dataclasses.dataclass
 class Configuration:
-  """A dataclass for configuring the behavior of `torchax`.
-
-  **Attributes:**
-
-  *   `debug_print_each_op` (`bool`): If `True`, prints each operation as it is
-      dispatched.
-  *   `debug_accuracy_for_each_op` (`bool`): If `True`, checks the accuracy of
-      each operation by comparing its output with the equivalent PyTorch
-      operation on the CPU.
-  *   `debug_mixed_tensor` (`bool`): If `True`, enables debugging for mixed
-      tensor operations.
-  *   `debug_print_each_op_operands` (`bool`): If `True`, prints the operands of
-      each operation.
-  *   `use_int32_for_index` (`bool`): If `True`, uses `int32` for indexing
-      operations.
-  *   `allow_mixed_math_with_scalar_tensor` (`bool`): If `True`, allows mixed
-      math operations between `torchax.Tensor` and scalar `torch.Tensor`s.
-  *   `force_materialize_views` (`bool`): If `True`, eagerly materializes `View`
-      objects into `torchax.Tensor`s.
-  *   `use_dlpack_for_data_conversion` (`bool`): If `True`, uses DLPack for
-      converting between `jax.Array` and `torch.Tensor`.
-  *   `use_tpu_flash_attention` (`bool`): If `True`, uses TPU-optimized flash
-      attention.
-  *   `shmap_flash_attention` (`bool`): If `True`, uses `shard_map` for flash
-      attention.
-  *   `treat_cuda_as_jax_device` (`bool`): If `True`, treats CUDA devices as JAX
-      devices.
-  *   `internal_respect_torch_return_dtypes` (`bool`): If `True`, respects the
-      return data types of PyTorch operations.
-  """
   debug_print_each_op: bool = False
   debug_accuracy_for_each_op: bool = False
   debug_mixed_tensor: bool = False

torchax/torchax/decompositions.py

@@ -1,10 +1,10 @@
-"""This file contains PyTorch operator decompositions that are not available in
-the stable version of PyTorch.
+"""This file contains some decompositons that are not available in torch stable.
 
-The decompositions are primarily sourced from the `main` branch of the PyTorch
-repository and are included here to provide support for newer operators. This
-module can also contain decompositions of a PyTorch op in terms of other
-PyTorch ops.
+Most likely from Content of
+https://github.com/pytorch/pytorch/blob/main/torch/_decomp/decompositions.py
+at main branch HEAD that we find useful here.
+
+Can also contain decompositions of a torch op in terms of other torch ops.
 """
 
 import functools
@@ -104,21 +104,18 @@ def _reflection_or_replication_pad(
 
 
 def bernoulli(self, *, generator=None):
-  """Decomposition for the `bernoulli` operator."""
   return (torch.rand_like(self, dtype=torch.float32) < self).to(self.dtype)
 
 
 _try_register(aten.bernoulli.default, bernoulli)
 
 
 def rand_like(self, **kwargs):
-  """Decomposition for the `rand_like` operator."""
   dtype = kwargs.get("dtype", self.dtype)
   return torch.rand(self.shape, dtype=dtype)
 
 
 def channel_shuffle(self, groups):
-  """Decomposition for the `channel_shuffle` operator."""
   batchsize, channels, height, width = self.shape
   channels_per_group = channels // groups
   self = self.reshape(batchsize, groups, channels_per_group, height, width)
@@ -134,7 +131,6 @@ def channel_shuffle(self, groups):
 
 
 def bernoulli_float(self, p=0.5):
-  """Decomposition for the `bernoulli_` operator with a float probability."""
   return self.bernoulli_(p)
 
 
@@ -154,10 +150,9 @@ def _grid_sampler_3d(
     padding_mode: int = 0,
     align_corners: bool = False,
 ) -> Tensor:
-  """Decomposition for the `grid_sampler_3d` operator.
+  """References: https://github.com/pytorch/pytorch/blob/06a7dc21c1005750598c37f3adbc031183c74de6/torch/_decomp/decompositions.py#L4075
 
-  This implementation is based on the 2D version in the PyTorch repository:
-  https://github.com/pytorch/pytorch/blob/06a7dc21c1005750598c37f3adbc031183c74de6/torch/_decomp/decompositions.py#L4075
+  The above implement the 2d case.
   """
   _expand_grid = False
   torch._check(
@@ -778,4 +773,4 @@ def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor,
 MUTABLE_DECOMPOSITION = [
     torch.ops.aten.bernoulli_.Tensor,
     torch.ops.aten.bernoulli_.float,
-]
+]

torchax/torchax/device_module.py

@@ -2,40 +2,32 @@
 
 
 def _is_in_bad_fork():
-  """Returns `False` as forking is not applicable in the same way as CUDA."""
   return False
 
 
 def manual_seed_all(seed):
-  """A placeholder for API compatibility; does not affect JAX's PRNG."""
   pass
 
 
 def device_count():
-  """Returns `1` as JAX manages devices as a single logical device."""
   return 1
 
 
 def get_rng_state():
-  """Returns an empty list for API compatibility."""
   return []
 
 
 def set_rng_state(new_state, device):
-  """A placeholder for API compatibility; does not affect JAX's PRNG."""
   pass
 
 
 def is_available():
-  """Returns `True` if JAX is available."""
   return True
 
 
 def current_device():
-  """Returns `0` as JAX manages devices as a single logical device."""
   return 0
 
 
 def get_amp_supported_dtype():
-  """Returns the data types supported by AMP (Automatic Mixed Precision)."""
   return [torch.float16, torch.bfloat16]