torchax 0.0.10.dev20251117__py3-none-any.whl

torchax/train.py

@@ -0,0 +1,132 @@
+# Copyright 2025 Google LLC
+#
+# 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.
+
+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,26 @@
+# Copyright 2025 Google LLC
+#
+# 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.
+
+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/util.py

@@ -0,0 +1,102 @@
+# Copyright 2025 Google LLC
+#
+# 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.
+
+from typing import Any, Callable
+
+
+def partition(original: list[Any],
+              func: Callable[[Any], bool]) -> tuple[list[Any], list[Any]]:
+  """Partitions elements into two parallel lists based on a predicate function.
+
+  Iterates through the 'original' list, applying 'func' to each element 'a'.
+  - If `func(a)` returns True, 'a' is appended to the first list ('truthy')
+    and `None` is appended to the second list ('falsy').
+  - If `func(a)` returns False, `None` is appended to the first list ('truthy')
+    and 'a' is appended to the second list ('falsy').
+
+  The result is two lists of the same length as the 'original' list, acting
+  as parallel representations of the partitioned elements, using `None` as
+  placeholders.
+
+  This is useful when we want to mark a group of elements as static (via passing
+  static_argnums) or donated (via donate_argnums) when combining with jax.jit
+  and friends.
+
+  Args:
+      original: The list of elements to partition.
+      func: A callable (function or lambda) that accepts an element from
+            'original' and returns a boolean value (True or False).
+
+  Returns:
+      A tuple containing two lists (`truthy`, `falsy`), both of the same
+      length as `original`:
+      - The first list contains elements `x` where `func(x)` was True, and
+        `None` otherwise.
+      - The second list contains elements `x` where `func(x)` was False, and
+        `None` otherwise.
+
+  Example:
+      >>> def is_even(n): return n % 2 == 0
+      >>> nums = [1, 2, 3, 4, 5, 6]
+      >>> truthy_list, falsy_list = partition(nums, is_even)
+      >>> truthy_list
+      [None, 2, None, 4, None, 6]
+      >>> falsy_list
+      [1, None, 3, None, 5, None]
+  """
+  truthy = []
+  falsy = []
+  for a in original:
+    t, f = (a, None) if func(a) else (None, a)
+    truthy.append(t)
+    falsy.append(f)
+  return truthy, falsy
+
+
+def merge(list1: list[Any], list2: list[Any]) -> list[Any]:
+  """Merges two lists element-wise, prioritizing non-None elements from list1.
+
+  Creates a new list where each element is taken from the corresponding position
+  in 'list1', unless that element is None, in which case the element from the
+  corresponding position in 'list2' is used. Assumes both lists have the
+  same length.
+
+  Invariant: merge(*partion(input_list, predicate)) == input_list for any predicate
+
+  Args:
+      list1: The primary list. Its elements are preferred unless they are None.
+      list2: The secondary list. Its elements are used as fallbacks when the
+              corresponding element in list1 is None.
+
+  Returns:
+      A new list representing the merged result.
+
+  Raises:
+      AssertionError: If 'list1' and 'list2' do not have the same length.
+
+  Example:
+      >>> l1 = [1, None, 3, None]
+      >>> l2 = [None, 2, None, 4]
+      >>> merge(l1, l2)
+      [1, 2, 3, 4]
+      >>> l3 = [None, 'b', None]
+      >>> l4 = ['a', None, 'c']
+      >>> merge(l3, l4)
+      ['a', 'b', 'c']
+  """
+  assert len(list1) == len(list2)
+  res = []
+  for a, b in zip(list1, list2):
+    res.append(b if a is None else a)
+  return res

torchax/view.py

@@ -0,0 +1,391 @@
+# Copyright 2025 Google LLC
+#
+# 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.
+
+import torch
+import torch.utils._pytree as torch_pytree
+import jax
+from enum import Enum
+from typing import Union, List, Tuple, Optional, Any, cast
+from abc import ABC, abstractmethod
+
+# Reference to original PyTorch native functions
+# https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml
+
+
+class ViewInfoType(Enum):
+  INVALID = 0
+  NARROW = 1
+  NO_OP = 2
+  PERMUTE = 3
+  RESHAPE = 4
+  RESIZE = 5
+  SELECT = 6
+  AS_STRIDED = 7
+  DIAGONAL = 8
+
+
+class ViewInfo(ABC):
+  """
+    Abstract base class for all view operations.
+    Defines the interface for applying and updating view transformations.
+    """
+
+  def __init__(
+      self,
+      view_info_type: ViewInfoType = ViewInfoType.INVALID,
+  ):
+    """
+        Initialize a ViewInfo object.
+
+        Args:
+            view_info_type: The type of view operation
+        """
+    self.view_info_type = view_info_type
+
+  @abstractmethod
+  def update_tensor(self, new_value: jax.Array,
+                    jax_array: jax.Array) -> jax.Array:
+    """
+        Apply this view transformation to a JAX array and update its value.
+
+        Args:
+            new_value: The new values to set in the view
+            jax_array: The parent array to update
+
+        Returns:
+            Updated array
+        """
+    pass
+
+  @abstractmethod
+  def transform_tensor(self, jax_array: jax.Array) -> jax.Array:
+    """
+        Apply this view transformation to a JAX array.
+
+        Args:
+            jax_array: The array to transform
+
+        Returns:
+            Transformed array
+        """
+    pass
+
+  @abstractmethod
+  def calculate_output_shape(self, source: jax.Array) -> List[int]:
+    """
+        Calculate the resulting shape after applying this view.
+
+        Args:
+            source: Original jax array before transformation
+
+        Returns:
+            Resulting shape after transformation
+        """
+    pass
+
+
+class NarrowInfo(ViewInfo):
+  """
+    Represents a slicing operation on a tensor.
+    Handles operations like tensor[1:3, :, 2:5:2].
+    """
+
+  def __init__(self, slices: Union[slice, Tuple[slice]]) -> None:
+    """
+        Args:
+            slices: The slice(s) to apply to the tensor.
+                E.g. jax_array.at[slices] will return the transformed tensor.
+        """
+    super().__init__(ViewInfoType.NARROW)
+    self.slices = slices
+
+  def __eq__(self, other: object) -> bool:
+    if not isinstance(other, NarrowInfo):
+      return False
+    return self.slices == other.slices
+
+  def transform_tensor(self, jax_array: jax.Array) -> jax.Array:
+    try:
+      return jax_array[self.slices]
+    except IndexError as e:
+      raise IndexError("Invalid slice operation") from e
+
+  def update_tensor(self, new_value: jax.Array,
+                    jax_array: jax.Array) -> jax.Array:
+    return jax_array.at[self.slices].set(new_value)
+
+  def calculate_output_shape(self, source: jax.Array) -> List[int]:
+    return source[self.slices].shape
+
+
+class SelectInfo(ViewInfo):
+  """
+    Represents a selection operation on a tensor.
+    Typically used for indexing operations that select specific elements.
+    """
+
+  def __init__(self,
+               dim: int = 0,
+               start: int = 0,
+               end: int = 0,
+               stride: int = 0) -> None:
+    super().__init__(ViewInfoType.SELECT)
+    self.dim: int = dim
+    self.start: int = start
+    self.end: int = end
+    self.stride: int = stride
+
+  def __eq__(self, other: object) -> bool:
+    if not isinstance(other, SelectInfo):
+      return False
+    return (self.dim == other.dim and self.start == other.start and
+            self.end == other.end and self.stride == other.stride)
+
+  def transform_tensor(self, jax_array: jax.Array) -> jax.Array:
+    raise NotImplementedError("SelectInfo.apply not implemented")
+
+  def update_tensor(self, new_value: jax.Array,
+                    jax_array: jax.Array) -> jax.Array:
+    raise NotImplementedError("SelectInfo.update not implemented")
+
+  def calculate_output_shape(self, source: jax.Array) -> List[int]:
+    raise NotImplementedError(
+        "SelectInfo.calculate_output_shape not implemented")
+
+
+class AsStridedInfo(ViewInfo):
+  """
+    Information for as_strided operations.
+    """
+
+  def __init__(self, stride: List[int], offset: int = 0) -> None:
+    super().__init__(ViewInfoType.AS_STRIDED)
+    self.stride: List[int] = stride
+    self.offset: int = offset
+
+  def __eq__(self, other: object) -> bool:
+    if not isinstance(other, AsStridedInfo):
+      return False
+    return self.offset == other.offset and self.stride == other.stride
+
+  def transform_tensor(self, jax_array: jax.Array) -> jax.Array:
+    raise NotImplementedError("AsStridedInfo.apply not implemented")
+
+  def update_tensor(self, new_value: jax.Array,
+                    jax_array: jax.Array) -> jax.Array:
+    raise NotImplementedError("AsStridedInfo.update not implemented")
+
+  def calculate_output_shape(self, source: jax.Array) -> List[int]:
+    raise NotImplementedError(
+        "AsStridedInfo.calculate_output_shape not implemented")
+
+
+class DiagonalInfo(ViewInfo):
+  """
+    Information for diagonal operations.
+    Extracts diagonal elements from a tensor.
+    """
+
+  def __init__(self, offset: int = 0, dim1: int = 0, dim2: int = 1) -> None:
+    """
+        Args:
+            offset: Offset from the main diagonal
+            dim1: First dimension for diagonal extraction
+            dim2: Second dimension for diagonal extraction
+        """
+    super().__init__(ViewInfoType.DIAGONAL)
+    self.offset: int = offset
+    self.dim1: int = dim1
+    self.dim2: int = dim2
+
+  def __eq__(self, other: object) -> bool:
+    if not isinstance(other, DiagonalInfo):
+      return False
+    return (self.offset == other.offset and self.dim1 == other.dim1 and
+            self.dim2 == other.dim2)
+
+  def transform_tensor(self, jax_array: jax.Array) -> jax.Array:
+    raise NotImplementedError("DiagonalInfo.apply not implemented")
+
+  def update_tensor(self, new_value: jax.Array,
+                    jax_array: jax.Array) -> jax.Array:
+    raise NotImplementedError("DiagonalInfo.update not implemented")
+
+  def calculate_output_shape(self, source: jax.Array) -> List[int]:
+    raise NotImplementedError(
+        "DiagonalInfo.calculate_output_shape not implemented")
+
+
+class View(torch.Tensor):
+  """
+    A View is a reference to another Tensor or another View,
+    with a transformation applied to it.
+    """
+
+  @staticmethod
+  def __new__(cls, parent: Union["torchax.Tensor", "View"], view_info: ViewInfo,
+              env: Any) -> "View":
+    """
+        Args:
+            parent: Parent tensor or view
+            view_info: Information about the view transformation
+            env: Environment for tensor operations
+        """
+    shape = view_info.calculate_output_shape(parent.jax())
+    return torch.Tensor._make_wrapper_subclass(
+        cls,
+        shape,
+        device="meta",
+        dtype=parent.dtype,
+        requires_grad=False,
+    )
+
+  def __init__(self, parent: Union["torchax.Tensor", "View"],
+               view_info: ViewInfo, env: Any) -> None:
+    super().__init__()
+    self.parent = parent
+    self.view_info = view_info
+    self._env = env
+
+  def get_transformation_chain(self) -> List[ViewInfo]:
+    """
+        Get all view transformations from the source tensor to this view.
+        """
+    if isinstance(self.parent, View):
+      transformations = self.parent.get_transformation_chain()
+      transformations.append(self.view_info)
+      return transformations
+    else:
+      return [self.view_info]
+
+  __torch_function__ = torch._C._disabled_torch_function_impl
+
+  def source_jax(self) -> jax.Array:
+    """
+        Returns the source tensor.
+        """
+    if isinstance(self.parent, View):
+      return self.parent.source_jax()
+    else:
+      return self.parent.jax()
+
+  def replace_source_jax(self, new_value: jax.Array) -> None:
+    """
+        Update the source tensor with new values.
+        """
+    if isinstance(self.parent, View):
+      self.parent.replace_source_jax(new_value)
+    else:
+      assert new_value.shape == self.parent._elem.shape
+      self.parent._elem = new_value
+
+  def torch(self) -> "torchax.Tensor":
+    """
+        Returns a Torchax tensor representing this view after all transformations
+        """
+    from torchax.tensor import Tensor
+
+    return Tensor(self.jax(), self._env)
+
+  def update(
+      self,
+      new_values: Union[jax.Array, "View", "torchax.Tensor"],
+      view_infos: Optional[List[ViewInfo]] = None,
+  ) -> None:
+    """
+        Update this view with new values, propagating changes back to source.
+        If view_infos is None, it will use the transformation chain
+        from the source tensor.
+        """
+    if view_infos is None:
+      view_infos = self.get_transformation_chain()
+
+    # Get the source JAX array
+    source_array = self.source_jax()
+
+    # Get the new value
+    from torchax.tensor import Tensor
+
+    if isinstance(new_values, View) or isinstance(new_values, Tensor):
+      new_values = new_values.jax()
+
+    # Apply all view transformations to the source array
+    # And store intermediate values
+    intermediate_values = [source_array]
+    for view_info in view_infos[:-1]:
+      intermediate_values.append(
+          view_info.transform_tensor(intermediate_values[-1]))
+
+    # TODO: Investigate efficiency of this algorithm
+    # Update the source array with the new value by
+    # applying inverse transformations in reverse order
+    for view_info, parent_array in zip(
+        reversed(view_infos), reversed(intermediate_values)):
+      # Apply the inverse transformation to propagate changes back
+      new_values = view_info.update_tensor(new_values, parent_array)
+
+    # Update the source tensor with the new values
+    self.replace_source_jax(new_values)
+
+  @classmethod
+  def __torch_dispatch__(
+      cls,
+      func: Any,
+      types: Tuple[Any, ...],
+      args: Tuple[Any, ...] = (),
+      kwargs: Optional[dict] = None,
+  ) -> Any:
+    raise AssertionError(
+        'torchax Tensors can only do math within the torchax environment.'
+        'Please wrap your code with `with torchax.default_env()` or '
+        'call torchax.enable_globally() before.')
+
+  def create_sub_view(self, view_info: ViewInfo) -> "View":
+    """
+        Create a new view that is a child of this view.
+        """
+    return View(self, view_info, self._env)
+
+  def __str__(self) -> str:
+    return f"View({self.torch()})"
+
+  def jax(self) -> jax.Array:
+    """
+        Returns a copy of the source tensor after transformations.
+        """
+    result = self.source_jax()
+    for view_info in self.get_transformation_chain():
+      result = view_info.transform_tensor(result)
+    return result
+
+  def __setitem__(self, indexes, val):
+    view_infos = self.get_transformation_chain() + [NarrowInfo(indexes)]
+    self.update(view_infos=view_infos, new_values=val)
+
+  def dim(self):
+    return self.ndim
+
+  @property
+  def device(self):
+    return torch.device("jax:0")
+
+  @property
+  def jax_device(self):
+    return self.jax().device
+
+  @property
+  def ndim(self):
+    return len(self.shape)
+
+  __repr__ = __str__