torchax 0.0.10.dev20251118__py3-none-any.whl

torchax/interop.py

@@ -0,0 +1,369 @@
+# 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 copy
+import functools
+from functools import wraps
+from inspect import signature
+
+import jax
+import jax.numpy as jnp
+import torch
+from jax import tree_util as pytree
+from jax.experimental.shard_map import shard_map
+from torch.nn.utils import stateless as torch_stateless
+
+import torchax
+from torchax import tensor, util
+from torchax.ops import mappings
+from torchax.types import JaxCallable, JaxValue, TorchCallable, TorchValue
+
+
+def extract_all_buffers(m: torch.nn.Module):
+  buffers = {}
+  params = {}
+
+  def extract_one(module, prefix):
+    for k in dir(module):
+      try:
+        v = getattr(module, k)
+      except Exception:
+        continue
+      qual_name = prefix + k
+      if isinstance(v, torch.nn.parameter.Parameter) and v.requires_grad:
+        params[qual_name] = v
+      elif isinstance(v, torch.Tensor):
+        buffers[qual_name] = v
+    for name, child in module.named_children():
+      extract_one(child, prefix + name + ".")
+
+  extract_one(m, "")
+  return params, buffers
+
+
+def set_all_buffers(m, params, buffers):
+  def set_one(module, prefix):
+    for k in dir(module):
+      qual_name = prefix + k
+      if (potential_v := buffers.get(qual_name)) is not None:
+        setattr(module, k, potential_v)
+      elif (potential_v := params.get(qual_name)) is not None:
+        print(k, potential_v)
+        setattr(module, k, torch.nn.Parameter(potential_v))
+    for name, child in module.named_children():
+      set_one(child, prefix + name + ".")
+
+  set_one(m, "")
+
+
+class JittableModule(torch.nn.Module):
+  def __init__(self, m: torch.nn.Module, extra_jit_args=None, dedup_parameters=True):
+    if extra_jit_args is None:
+      extra_jit_args = {}
+    super().__init__()
+    self.params, self.buffers = extract_all_buffers(m)
+    self._model = m
+    self._jitted = {}
+
+    self._extra_jit_args = extra_jit_args
+
+    self._extra_dumped_weights = {}
+
+    if dedup_parameters:
+      temp = collections.defaultdict(list)
+      for k, v in self.params.items():
+        temp[id(v)].append(k)
+
+      for v in temp.values():
+        if len(v) > 1:
+          # duplicated weights with different name
+          self._extra_dumped_weights[v[0]] = v[1:]
+          for extra_keys in v[1:]:
+            del self.params[extra_keys]
+
+  @property
+  def __class__(self):
+    # Lie about the class type so that
+    # isinstance(jittable_module, self._model.__class__) works
+    return self._model.__class__
+
+  def __call__(self, *args, **kwargs):
+    return self.forward(*args, **kwargs)
+
+  def functional_call(self, method_or_name, params, buffers, *args, **kwargs):
+    kwargs = kwargs or {}
+    params_copy = copy.copy(params)
+    params_copy.update(buffers)
+    # reinflate the state dict so there are not any missing keys
+    for k, v in self._extra_dumped_weights.items():
+      for new_key in v:
+        params_copy[new_key] = params_copy[k]
+
+    if isinstance(method_or_name, str):
+      method = getattr(self._model, method_or_name)
+    else:
+      if not callable(method_or_name):
+        raise TypeError(
+          f"method_or_name should be a callable or a string, got {type(method_or_name)}"
+        )
+      method = method_or_name
+      args = (self._model,) + args
+    with torch_stateless._reparametrize_module(self._model, params_copy):
+      res = method(*args, **kwargs)
+    return res
+
+  def jittable_call(self, method_name: str, *args, **kwargs):
+    if method_name not in self._jitted:
+      jitted = jax_jit(
+        functools.partial(self.functional_call, method_name),
+        kwargs_for_jax_jit=self._extra_jit_args,
+      )
+
+      def jitted_forward(*args, **kwargs):
+        return jitted(self.params, self.buffers, *args, **kwargs)
+
+      self._jitted[method_name] = jitted_forward
+    return self._jitted[method_name](*args, **kwargs)
+
+  def forward(self, *args, **kwargs):
+    return self.jittable_call("forward", *args, **kwargs)
+
+  def __getattr__(self, key):
+    if key == "_model":
+      return super().__getattr__(key)
+    if key in self._jitted:
+      return self._jitted[key]
+    return getattr(self._model, key)
+
+  def make_jitted(self, key):
+    jitted = jax_jit(
+      functools.partial(self.functional_call, key),
+      kwargs_for_jax_jit=self._extra_jit_args,
+    )
+
+    def call(*args, **kwargs):
+      return jitted(self.params, self.buffers, *args, **kwargs)
+
+    self._jitted[key] = call
+
+
+class CompileMixin:
+  def functional_call(self, method, params, buffers, *args, **kwargs):
+    kwargs = kwargs or {}
+    params_copy = copy.copy(params)
+    params_copy.update(buffers)
+    with torch_stateless._reparametrize_module(self, params_copy):
+      res = method(*args, **kwargs)
+    return res
+
+  def jit(self, method):
+    jitted = jax_jit(functools.partial(self.functional_call, method_name))  # noqa: F821
+
+    def call(*args, **kwargs):
+      return jitted(self.named_paramters(), self.named_buffers(), *args, **kwargs)
+
+    return call
+
+
+def compile_nn_module(m: torch.nn.Module, methods=None):
+  if methods is None:
+    methods = ["forward"]
+
+  type(
+    m.__class__.__name__ + "_with_CompileMixin",
+    (CompileMixin, m.__class__),
+  )
+  m.__class__ = NewParent  # noqa: F821
+
+
+def _torch_view(t: JaxValue) -> TorchValue:
+  # t is an object from jax land
+  # view it as-if it's a torch land object
+  if isinstance(t, jax.Array):
+    return tensor.Tensor(t, torchax.default_env())
+  if isinstance(t, jnp.dtype):
+    return mappings.j2t_dtype(t)
+  if callable(t):  # t is a JaxCallable
+    return functools.partial(call_jax, t)
+  # regular types are not changed
+  return t
+
+
+torch_view = functools.partial(pytree.tree_map, _torch_view)
+
+
+def _jax_view(t: TorchValue) -> JaxValue:
+  # t is an object from torch land
+  # view it as-if it's a jax land object
+  if isinstance(t, torch.Tensor):
+    assert isinstance(t, tensor.Tensor) or isinstance(t, tensor.View), type(t)
+    return t.jax()
+  if isinstance(t, type(torch.int32)):
+    return mappings.t2j_dtype(t)
+
+  # torch.nn.Module needs special handling
+  if not isinstance(t, torch.nn.Module) and callable(t):  # t is a TorchCallable
+    return functools.partial(call_torch, t)
+  # regular types are not changed
+  return t
+
+
+jax_view = functools.partial(pytree.tree_map, _jax_view)
+
+
+def call_jax(
+  jax_func: JaxCallable, *args: TorchValue, **kwargs: TorchValue
+) -> TorchValue:
+  args, kwargs = jax_view((args, kwargs))
+  res: JaxValue = jax_func(*args, **kwargs)
+  return torch_view(res)
+
+
+def call_torch(
+  torch_func: TorchCallable, *args: JaxValue, **kwargs: JaxValue
+) -> JaxValue:
+  args, kwargs = torch_view((args, kwargs))
+  with torchax.default_env():
+    res: TorchValue = torch_func(*args, **kwargs)
+  return jax_view(res)
+
+
+def j2t_autograd(fn, call_jax=call_jax):
+  """Given a JAX function, returns a PyTorch autograd function implemented with `jax.vjp(fn)`.
+
+  It wraps `fn` with `jax.vjp` to compute both the output and residuals (intermediate
+  activations). The wrapped function is then run via `call_jax` and integrated into
+  the PyTorch autograd framework by saving the residuals into the context object.
+  """
+
+  # NOTE(qihqi): This function cannot be inlined from the callsite
+  #  Becuase if it does, then it won't hit the compilation cache for
+  #  call_jax. Call jax uses functions' id as key.
+  # It is nested inside j2t_autograd to ensure it gets a unique ID for each
+  # wrapped pure function, preventing cache collisions between different pure modules.
+  def _jax_forward(fn, other, tree_def, tensors):
+    """JAX function to compute output and vjp function.
+
+    primals should be a tuple (args, kwargs).
+    """
+    import jax
+    from jax.tree_util import tree_unflatten
+
+    def fn_wrapper(*tensors):
+      # Reconstruct the original args and kwargs
+      flat_inputs = util.merge(tensors, other)
+      args, kwargs = tree_unflatten(tree_def, flat_inputs)
+      return fn(*args, **kwargs)
+
+    return jax.vjp(fn_wrapper, *tensors)
+
+  def _jax_backward(vjp_spec, saved_tensors, grad_out):
+    """JAX function to compute input gradients.
+
+    Unflattening `saved_tensors` with `vjp_spec` should restore the original vjp function.
+    """
+    from jax.tree_util import tree_unflatten
+
+    fun_vjp = tree_unflatten(vjp_spec, saved_tensors)
+    return fun_vjp(grad_out)
+
+  @wraps(fn)
+  def inner(*args, **kwargs):
+    from jax.tree_util import tree_flatten
+
+    class JaxFun(torch.autograd.Function):
+      @staticmethod
+      def forward(ctx, tree_def, *flat_args_kwargs):
+        tensors, other = util.partition(
+          flat_args_kwargs, lambda x: isinstance(x, torch.Tensor)
+        )
+        # We want the arguments that don't require grads to be closured?
+
+        y, fun_vjp = call_jax(_jax_forward, fn, other, tree_def, tensors)
+
+        # Save necessary information for backward
+        # Flatten the vjp function. `vjp_spec` contains a jaxpr for the backward pass.
+        # `residuals` contains the tensors needed for the backward pass.`
+        residuals, vjp_spec = tree_flatten(fun_vjp)
+        ctx.vjp_spec = vjp_spec
+        ctx.save_for_backward(*residuals)
+        return y
+
+      @staticmethod
+      def backward(ctx, *grad_out):
+        assert len(grad_out) > 0
+        grad_out = grad_out if len(grad_out) > 1 else grad_out[0]
+
+        input_grads_structured = call_jax(
+          _jax_backward, ctx.vjp_spec, ctx.saved_tensors, grad_out
+        )
+
+        # Construct the gradient tuple to be returned.
+        # It needs to match the inputs to forward: (tree_def, *flat_inputs)
+        # The first gradient (for tree_def) is None.
+        # The subsequent gradients correspond to flat_inputs.
+        # We need to put a None for inputs that did not require gradients.
+        final_grads = [None]
+        for needs_grad, grad in zip(
+          ctx.needs_input_grad[1:], input_grads_structured, strict=True
+        ):
+          final_grads.append(grad if needs_grad else None)
+
+        return tuple(final_grads)
+
+    sig = signature(fn)
+    bound = sig.bind(*args, **kwargs)
+    bound.apply_defaults()
+    flat_args_kwargs, tree_def = tree_flatten((bound.args, bound.kwargs))
+    y = JaxFun.apply(tree_def, *flat_args_kwargs)
+    return y
+
+  return inner
+
+
+fori_loop = torch_view(jax.lax.fori_loop)
+
+
+def wrap_jax_jit(torch_function, jax_jit_func=jax.jit, kwargs_for_jax=None):
+  kwargs_for_jax = kwargs_for_jax or {}
+  jax_func = jax_view(torch_function)
+  jitted = jax_jit_func(jax_func, **kwargs_for_jax)
+  return torch_view(jitted)
+
+
+def jax_jit(torch_function, kwargs_for_jax_jit=None, fix_for_buffer_donation=False):
+  return wrap_jax_jit(
+    torch_function, jax_jit_func=jax.jit, kwargs_for_jax=kwargs_for_jax_jit
+  )
+
+
+def jax_shard_map(torch_function, kwargs_for_jax_shard_map=None):
+  return wrap_jax_jit(
+    torch_function, jax_jit_func=shard_map, kwargs_for_jax=kwargs_for_jax_shard_map
+  )
+
+
+def jax_value_and_grad(torch_function, kwargs_for_value_and_grad=None):
+  return wrap_jax_jit(
+    torch_function,
+    jax_jit_func=jax.value_and_grad,
+    kwargs_for_jax=kwargs_for_value_and_grad,
+  )
+
+
+def gradient_checkpoint(torch_function, kwargs=None):
+  return wrap_jax_jit(
+    torch_function, jax_jit_func=jax.checkpoint, kwargs_for_jax=kwargs
+  )

torchax/mesh_util.py

@@ -0,0 +1,236 @@
+# 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 jax
+import torch
+from jax.sharding import NamedSharding, PartitionSpec
+
+import torchax
+from torchax import interop
+
+
+def _shard_first_multiple_of(axis_name, shape, multiple_of):
+  """Creates a PartitionSpec to shard the first dimension divisible by a number.
+
+  Iterates through the dimensions specified by `shape`. Finds the first dimension
+  whose size is a multiple of `multiple_of` and returns a PartitionSpec that
+  shards that dimension along the given `axis_name`. All preceding dimensions
+  are not sharded (marked as None in the PartitionSpec). All subsequent dimensions
+  skipped, which would be implicitly treated as replicated.
+
+  Args:
+    axis_name: The name of the mesh axis to shard along (e.g., "data", "mdl").
+    shape: A tuple or list representing the shape of the tensor to be sharded.
+    multiple_of: The integer value that a dimension size must be divisible by
+      in order to be sharded. Typically the size of the mesh axis.
+
+  Returns:
+    A jax.sharding.PartitionSpec object specifying how to shard the tensor.
+    For example, if shape=(10, 20, 30), axis_name='x', multiple_of=4,
+    it would return PartitionSpec(None, 'x', None).
+    If none divides then it should return a replicated PartitionSpec
+  """
+  sharding = []
+  found = False
+  for size in shape:
+    if not found and size % multiple_of == 0:
+      found = True
+      sharding.append(axis_name)
+    else:
+      sharding.append(None)
+  return PartitionSpec(*sharding)
+
+
+class SingleAxisSharder:
+  """A callable object that generates PartitionSpecs for single-axis sharding.
+
+  This sharder strategy attempts to shard the *first* dimension of a tensor
+  that is divisible by the specified `axis_size` along the given `axis_name`.
+  It's useful for simple 1D mesh sharding scenarios like FSDP where parameters
+  are typically sharded along one dimension.
+
+  Attributes:
+    axis_name: The name of the mesh axis to shard along.
+    axis_size: The size of the mesh axis (number of devices along that axis).
+  """
+
+  def __init__(self, axis_name, axis_size, replicate_unshardable=False):
+    """Initializes the SingleAxisSharder.
+
+    Args:
+      axis_name: The name of the mesh axis (e.g., "fsdp", "data").
+      axis_size: The number of devices along the specified mesh axis.
+      replicate_unshardable: indicate whether it should return replicated sharding
+        (P()) when none of the axis is divisible by the axis size.
+    """
+    self.axis_name = axis_name
+    self.axis_size = axis_size
+    self.replicate_unshardable = replicate_unshardable
+
+  def __call__(self, name, shapedtype):
+    """Generates a PartitionSpec for a given tensor name and shaped type.
+
+    Args:
+      name: The name of the tensor (e.g., parameter name). This argument is
+        provided for compatibility with more complex sharders but is not used
+        by this simple sharder.
+      shapedtype: An object with a `.shape` attribute describing the tensor's shape,
+        and `.dtype` describing it's dtype. Example: jax.Array, jax.ShapeDtypeStruct
+        or a torch.Tensor)
+
+    Returns:
+      A jax.sharding.PartitionSpec determined by finding the first dimension
+      in `shapedtype.shape` divisible by `self.axis_size` using the helper
+      `_shard_first_multiple_of`.
+    """
+    del name
+    sharding = _shard_first_multiple_of(
+      self.axis_name, shapedtype.shape, self.axis_size
+    )
+    if not self.replicate_unshardable and all(s is None for s in sharding):
+      raise AssertionError(
+        f"Unable to find a dim to shard because "
+        f"None of the dims ({shapedtype.shape}) in shape is multiple of {self.axis_size}"
+      )
+    return sharding
+
+
+class Mesh:
+  """A helper class that wraps `jax.sharding.Mesh` object.
+
+  The goal of this class is to provide helper methods that facilitate the
+  sharding of PyTorch tensors or models given a JAX device mesh configuration.
+  It simplifies initializing models directly into a sharded state.
+
+  Attributes:
+    jax_mesh: The underlying `jax.sharding.Mesh` object defining the device grid
+      and axis names.
+    _sharder: The default sharding strategy callable (like SingleAxisSharder)
+      used to determine the PartitionSpec for each parameter if not overridden
+      during method calls. Can be None if no default is appropriate or set.
+  """
+
+  @classmethod
+  def fsdp_mesh(cls, axis_name="fsdp"):
+    """Creates a Mesh instance suitable for 1D FSDP-style sharding.
+
+    This named constructor creates a 1D mesh encompassing all available XLA
+    devices. It assigns the specified `axis_name` to this single dimension.
+    It then creates a `Mesh` instance using this JAX mesh and a
+    `SingleAxisSharder` configured appropriately for this 1D mesh.
+
+    Args:
+      axis_name: The name to assign to the single mesh axis (default: "fsdp").
+        This name will be used by the default `SingleAxisSharder`.
+
+    Returns:
+      A Mesh instance configured with a 1D JAX mesh across all devices and a
+      corresponding SingleAxisSharder.
+    """
+    ndevice = jax.device_count()
+    jax_mesh = jax.make_mesh((ndevice,), (axis_name,))
+    # replicate_unshardable so scalars and small model attributes are replicated.
+    return cls(jax_mesh, SingleAxisSharder(axis_name, ndevice, True))
+
+  def __init__(self, jax_mesh, sharder=None):
+    """Initializes the Mesh helper.
+
+    Args:
+      jax_mesh: A pre-configured `jax.sharding.Mesh` object defining the
+        physical device grid and logical axis names.
+      sharder: An optional callable (e.g., an instance of SingleAxisSharder)
+        that takes (name, shapedtype) and returns a `jax.sharding.PartitionSpec`.
+        This serves as the default sharding strategy.
+        If None, and the provided `jax_mesh` has exactly one axis, a
+        `SingleAxisSharder` is created automatically for that single axis.
+        If None and the mesh has multiple axes, `_sharder` remains None, and
+        an `override_sharder` must be provided to methods like
+        `initialize_model_sharded`.
+    """
+    self.jax_mesh = jax_mesh
+    if sharder is None:
+      assert len(self.jax_mesh.axis_names) == 1
+      sharder = SingleAxisSharder(
+        self.jax_mesh.axis_names[0], len(self.mesh.device_ids)
+      )
+    self._sharder = sharder
+
+  def initialize_model_sharded(
+    self, model_class, init_args, init_kwargs=None, override_sharder=None
+  ):
+    """Initializes a PyTorch model with its parameters sharded across the mesh.
+
+    This method orchestrates the initialization of a `torch.nn.Module` such
+    that its parameters are created directly on the target devices according
+    to the sharding specifications derived from the mesh and the chosen sharder.
+    It leverages `torchax.interop.jax_jit` to achieve this.
+
+    Args:
+      model_class: The PyTorch model class (a subclass of `torch.nn.Module`).
+      init_args: A tuple containing the positional arguments required by the
+        `model_class.__init__` method.
+      init_kwargs: An optional dictionary containing the keyword arguments for
+        the `model_class.__init__` method. Defaults to None (treated as {}).
+      override_sharder: An optional callable sharding strategy to use
+        specifically for this initialization. If provided, it takes precedence
+        over the mesh's default `_sharder`. It must accept `(name, shapedtype)`
+        and return a `PartitionSpec`. If None, the mesh's default `_sharder`
+        is used.
+
+    Returns:
+      An instance of `model_class` whose parameters have been initialized and
+      are represented by sharded tensors distributed across the devices in the
+      `jax_mesh`.
+
+    Raises:
+      ValueError: If no sharder is available (i.e., `override_sharder` is None
+        and the mesh's default `_sharder` is also None).
+      AssertionError: Can be raised by the sharder (e.g., `SingleAxisSharder`)
+        if it fails to determine a valid sharding for any parameter.
+      TypeError: If `shapedtype` passed to the sharder doesn't have a `.shape`.
+      Other errors from JAX JIT compilation or PyTorch model initialization.
+    """
+    init_kwargs = init_kwargs or {}
+    with torch.device("meta"), torchax.disable_temporarily():
+      model = model_class(*init_args, **init_kwargs)
+
+    sharder = override_sharder or self._sharder
+
+    states = model.state_dict()
+    output_shards = {
+      name: NamedSharding(self.jax_mesh, sharder(name, tensor))
+      for name, tensor in states.items()
+    }
+
+    def model_initializer():
+      with torchax.default_env(), torch.device("meta"):
+        model = model_class(*init_args, **init_kwargs)
+      return dict(model.state_dict())
+
+    jitted = interop.jax_jit(
+      model_initializer, kwargs_for_jax_jit={"out_shardings": output_shards}
+    )
+    weights_dict = jitted()
+
+    model.load_state_dict(weights_dict, assign=True)
+    return model
+
+  def shard_model(self, model, override_sharder=None):
+    sharder = override_sharder or self._sharder
+    states = model.state_dict()
+    output_shards = {
+      name: NamedSharding(self.jax_mesh, sharder(name, tensor))
+      for name, tensor in states.items()
+    }
+    model.load_state_dict(output_shards, assign=True)

torchax/ops/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+
+
+def all_aten_jax_ops():
+  # to load the ops
+  import torchax.ops.jaten  # type: ignore
+  import torchax.ops.ops_registry  # type: ignore
+
+  return {
+    key: val.func
+    for key, val in torchax.ops.ops_registry.all_aten_ops.items()
+    if val.is_jax_function
+  }