torchax 0.0.10.dev20251117__py3-none-any.whl

torchax/ops/jax_reimplement.py

@@ -0,0 +1,185 @@
+# 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 collections.abc import Sequence
+from jax._src.numpy.util import promote_dtypes_inexact
+import numpy as np
+import jax
+from jax import numpy as jnp
+from jax._src.util import canonicalize_axis
+from jax._src import core
+from jax._src.image.scale import _kernels, ResizeMethod
+from jax import lax
+from typing import Callable
+
+# TODO: This block of code needs to be revisited based on https://github.com/jax-ml/jax/issues/24106
+# START ----------------- JAX code copied for fixing scale_and_translate -----------------------------
+
+# JAX Link: https://github.com/jax-ml/jax/blob/18f48bd52abe907ff9818da52f3d195d32910c1b/jax/_src/image/scale.py#L52
+
+
+def compute_weight_mat(input_size: core.DimSize, output_size: core.DimSize,
+                       scale, translation, kernel: Callable, antialias: bool):
+  dtype = jnp.result_type(scale, translation)
+  inv_scale = 1. / scale
+  # When downsampling the kernel should be scaled since we want to low pass
+  # filter and interpolate, but when upsampling it should not be since we only
+  # want to interpolate.
+  kernel_scale = jnp.maximum(inv_scale, 1.) if antialias else 1.
+  sample_f = ((jnp.arange(output_size, dtype=dtype) + 0.5) * inv_scale -
+              translation * inv_scale - 0.5)
+  x = (
+      jnp.abs(sample_f[jnp.newaxis, :] -
+              jnp.arange(input_size, dtype=dtype)[:, jnp.newaxis]) /
+      kernel_scale)
+  weights = kernel(x)
+
+  total_weight_sum = jnp.sum(weights, axis=0, keepdims=True)
+  weights = jnp.where(
+      jnp.abs(total_weight_sum) > 1000. * float(np.finfo(np.float32).eps),
+      jnp.divide(weights, jnp.where(total_weight_sum != 0, total_weight_sum,
+                                    1)), 0)
+  # Zero out weights where the sample location is completely outside the input
+  # range.
+  # Note sample_f has already had the 0.5 removed, hence the weird range below.
+
+  # (barney-s) -------------- returning weights without zeroing ---------------------
+  return weights
+  input_size_minus_0_5 = core.dimension_as_value(input_size) - 0.5
+  return jnp.where(
+      jnp.logical_and(sample_f >= -0.5, sample_f
+                      <= input_size_minus_0_5)[jnp.newaxis, :], weights, 0)
+  # (barney-s) -------------- END returning weights without zeroing ---------------------
+
+
+# JAX Link: https://github.com/jax-ml/jax/blob/18f48bd52abe907ff9818da52f3d195d32910c1b/jax/_src/image/scale.py#L86
+
+
+def _scale_and_translate(x, output_shape: core.Shape,
+                         spatial_dims: Sequence[int], scale, translation,
+                         kernel, antialias: bool, precision):
+  input_shape = x.shape
+  assert len(input_shape) == len(output_shape)
+  assert len(spatial_dims) == len(scale)
+  assert len(spatial_dims) == len(translation)
+  if len(spatial_dims) == 0:
+    return x
+  contractions = []
+  in_indices = list(range(len(output_shape)))
+  out_indices = list(range(len(output_shape)))
+  for i, d in enumerate(spatial_dims):
+    d = canonicalize_axis(d, x.ndim)
+    m = input_shape[d]
+    n = output_shape[d]
+    w = compute_weight_mat(m, n, scale[i], translation[i], kernel,
+                           antialias).astype(x.dtype)
+    contractions.append(w)
+    contractions.append([d, len(output_shape) + i])
+    out_indices[d] = len(output_shape) + i
+  contractions.append(out_indices)
+  return jnp.einsum(x, in_indices, *contractions, precision=precision)
+
+
+# JAX Link: https://github.com/jax-ml/jax/blob/18f48bd52abe907ff9818da52f3d195d32910c1b/jax/_src/image/scale.py#L172
+
+
+# scale and translation here are scalar elements of an np.array, what is the
+# correct type annotation?
+def scale_and_translate(
+    image,
+    shape: core.Shape,
+    spatial_dims: Sequence[int],
+    scale,
+    translation,
+    # (barney-s) use string
+    method: str,  #(barney-s) | ResizeMethod,
+    antialias: bool = True,
+    precision=lax.Precision.HIGHEST):
+  """Apply a scale and translation to an image.
+
+  Generates a new image of shape 'shape' by resampling from the input image
+  using the sampling method corresponding to method. For 2D images, this
+  operation transforms a location in the input images, (x, y), to a location
+  in the output image according to::
+
+    (x * scale[1] + translation[1], y * scale[0] + translation[0])
+
+  (Note the *inverse* warp is used to generate the sample locations.)
+  Assumes half-centered pixels, i.e the pixel at integer location ``row, col``
+  has coordinates ``y, x = row + 0.5, col + 0.5``, and similarly for other input
+  image dimensions.
+
+  If an output location(pixel) maps to an input sample location that is outside
+  the input boundaries then the value for the output location will be set to
+  zero.
+
+  The ``method`` argument expects one of the following resize methods:
+
+  ``ResizeMethod.LINEAR``, ``"linear"``, ``"bilinear"``, ``"trilinear"``,
+    ``"triangle"`` `Linear interpolation`_. If ``antialias`` is ``True``, uses a
+    triangular filter when downsampling.
+
+  ``ResizeMethod.CUBIC``, ``"cubic"``, ``"bicubic"``, ``"tricubic"``
+    `Cubic interpolation`_, using the Keys cubic kernel.
+
+  ``ResizeMethod.LANCZOS3``, ``"lanczos3"``
+    `Lanczos resampling`_, using a kernel of radius 3.
+
+  ``ResizeMethod.LANCZOS5``, ``"lanczos5"``
+    `Lanczos resampling`_, using a kernel of radius 5.
+
+  .. _Linear interpolation: https://en.wikipedia.org/wiki/Bilinear_interpolation
+  .. _Cubic interpolation: https://en.wikipedia.org/wiki/Bicubic_interpolation
+  .. _Lanczos resampling: https://en.wikipedia.org/wiki/Lanczos_resampling
+
+  Args:
+    image: a JAX array.
+    shape: the output shape, as a sequence of integers with length equal to the
+      number of dimensions of `image`.
+    spatial_dims: A length K tuple specifying the spatial dimensions that the
+      passed scale and translation should be applied to.
+    scale: A [K] array with the same number of dimensions as image, containing
+      the scale to apply in each dimension.
+    translation: A [K] array with the same number of dimensions as image,
+      containing the translation to apply in each dimension.
+    method: the resizing method to use; either a ``ResizeMethod`` instance or a
+      string. Available methods are: LINEAR, LANCZOS3, LANCZOS5, CUBIC.
+    antialias: Should an antialiasing filter be used when downsampling? Defaults
+      to ``True``. Has no effect when upsampling.
+
+  Returns:
+    The scale and translated image.
+  """
+  shape = core.canonicalize_shape(shape)
+  if len(shape) != image.ndim:
+    msg = ('shape must have length equal to the number of dimensions of x; '
+           f' {shape} vs {image.shape}')
+    raise ValueError(msg)
+  if isinstance(method, str):
+    method = ResizeMethod.from_string(method)
+  if method == ResizeMethod.NEAREST:
+    # Nearest neighbor is currently special-cased for straight resize, so skip
+    # for now.
+    raise ValueError('Nearest neighbor resampling is not currently supported '
+                     'for scale_and_translate.')
+  assert isinstance(method, ResizeMethod)
+
+  kernel = _kernels[method]
+  image, = promote_dtypes_inexact(image)
+  scale, translation = promote_dtypes_inexact(scale, translation)
+  return _scale_and_translate(image, shape, spatial_dims, scale, translation,
+                              kernel, antialias, precision)
+
+
+# END ----------------- END JAX code copied for testing -----------------------------

torchax/ops/jc10d.py

@@ -0,0 +1,66 @@
+# 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 jax
+import jax.numpy as jnp
+
+from torchax.ops import ops_registry
+
+
+def op(*aten, **kwargs):
+
+  def inner(func):
+    for a in aten:
+      ops_registry.register_torch_dispatch_op(a, func, **kwargs)
+    return func
+
+  return inner
+
+
+@op(torch.ops._c10d_functional.all_gather_into_tensor)
+def _c10d_all_gather(input, group_size: int, group_name: str):
+  return jax.lax.all_gather(input, "torch_dist")
+
+
+@op(torch.ops._c10d_functional.all_reduce)
+def _c10d_all_reduce(self, reduceOp: str, group_name: str):
+
+  if reduceOp == "sum":
+    res = jax.lax.psum(self, axis_name="torch_dist")
+  elif reduceOp == "avg":
+    res = jax.lax.pmean(self, axis_name="torch_dist")
+  elif reduceOp == "min":
+    res = jax.lax.pmin(self, axis_name="torch_dist")
+  elif reduceOp == "max":
+    res = jax.lax.pmax(self, axis_name="torch_dist")
+  else:
+    raise RuntimeError(f"Reduce op {reduceOp} not implemented")
+  return res
+
+
+@op(torch.ops._c10d_functional.broadcast)
+def _c10d_broadcast(self, src: int, group_name: str):
+  masked = jnp.where(
+      jax.lax.axis_index("torch_dist") == src,
+      self,
+      jnp.zeros_like(self),
+  )
+  return jax.lax.psum(masked, "torch_dist")
+
+
+@op(torch.ops._c10d_functional.wait_tensor)
+def _c10d_wait_tensor(tensor):
+  # Async tensor is aleady `wait`ed by dispatcher
+  return tensor

torchax/ops/jimage.py

@@ -0,0 +1,127 @@
+# 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 jax.numpy as jnp
+
+
+def cubic_kernel(x, a=-0.75):
+  """Cubic kernel with a = -0.75 (PyTorch-like Keys kernel)"""
+  absx = jnp.abs(x)
+  x2 = absx * absx
+  x3 = x2 * absx
+  cond1 = (absx <= 1)
+  cond2 = (absx > 1) & (absx < 2)
+  f1 = (a + 2) * x3 - (a + 3) * x2 + 1
+  f2 = a * x3 - 5 * a * x2 + 8 * a * absx - 4 * a
+  return jnp.where(cond1, f1, jnp.where(cond2, f2, 0.0))
+
+
+def compute_contribs(in_size,
+                     out_size,
+                     scale,
+                     support=2.0,
+                     align_corners=False,
+                     dtype=None):
+  if align_corners:
+    if out_size == 1:
+      in_coords = jnp.zeros((1,), dtype=dtype)
+    else:
+      in_coords = jnp.linspace(0, in_size - 1, out_size, dtype=dtype)
+  else:
+    out_coords = jnp.arange(out_size, dtype=dtype) + 0.5
+    in_coords = out_coords / scale - 0.5
+
+  left_idx = jnp.floor(in_coords).astype(jnp.int32) - 1
+  idxs = left_idx[:, None] + jnp.arange(4)
+
+  dx = in_coords[:, None] - idxs
+
+  weights = cubic_kernel(dx)
+
+  weights = weights / jnp.sum(weights, axis=1, keepdims=True)
+  return idxs, weights
+
+
+def gather_weights(img, idxs, axis):
+  """Safely gather with boundary handling"""
+  idxs = jnp.clip(idxs, 0, img.shape[axis] - 1)
+  return jnp.take(img, idxs, axis=axis)
+
+
+def interpolate_along_axis_bchw(img, idxs, weights, axis):
+  """
+    Interpolate along H (axis=2) or W (axis=3) for tensor (B, C, H, W).
+    idxs: (out_size, 4) int32 indices
+    weights: (out_size, 4) float32 weights
+    """
+  assert axis in (2, 3), "Axis must be 2 (H) or 3 (W)"
+  out_size = idxs.shape[0]
+  k = idxs.shape[1]  # Typically 4 for cubic
+
+  # Clip to input bounds
+  idxs = jnp.clip(idxs, 0, img.shape[axis] - 1)  # (out_size, 4)
+
+  def gather_and_weight(i):
+    idx = idxs[i]  # (4,)
+    w = weights[i]  # (4,)
+
+    def gather_one(offset):
+      return jnp.take(img, idx[offset], axis=axis)  # shape (B, C, H, W)
+
+    gathered = jnp.stack([gather_one(o) for o in range(k)],
+                         axis=0)  # (4, B, C, H, W)
+    weighted = jnp.tensordot(w, gathered, axes=(0, 0))  # (B, C, H, W)
+    return weighted
+
+  out = jax.vmap(gather_and_weight)(
+      jnp.arange(out_size))  # (out_size, B, C, H, W)
+
+  # Move the interpolated axis back into place
+  if axis == 2:  # interpolated over H
+    return jnp.moveaxis(out, 0, 2)  # (B, C, out_H, W)
+  else:  # axis == 3, interpolated over W
+    return jnp.moveaxis(out, 0, 3)  # (B, C, H, out_W)
+
+
+def interpolate_bicubic_no_aa(img, out_h, out_w, align_corners=False):
+  h, w = img.shape[-2:]
+  if align_corners and out_h > 1:
+    scale_y = (h - 1) / (out_h - 1)
+  else:
+    scale_y = out_h / h
+
+  if align_corners and out_w > 1:
+    scale_x = (w - 1) / (out_w - 1)
+  else:
+    scale_x = out_w / w
+
+  idxs_y, weights_y = compute_contribs(
+      h,
+      out_h,
+      scale_y,
+      align_corners=align_corners,
+      dtype=img.dtype,
+  )
+  tmp = interpolate_along_axis_bchw(img, idxs_y, weights_y, axis=2)
+
+  idxs_x, weights_x = compute_contribs(
+      w,
+      out_w,
+      scale_x,
+      align_corners=align_corners,
+      dtype=img.dtype,
+  )
+  out = interpolate_along_axis_bchw(tmp, idxs_x, weights_x, axis=3)
+  return out

torchax/ops/jlibrary.py

@@ -0,0 +1,94 @@
+# 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.
+
+"""The `jlibrary` module has functions which help to preserve torch.library ops
+during export. This includes aten ops, and custom operations.
+"""
+
+import torch
+import torch.nn as nn
+import torchax
+from torchax.ops import jaten
+import jax
+import functools
+
+
+def _jit_composite_impl(composite_name, jaxpr_impl, **jit_args):
+  """Wrap a jaxpr in a jitted function with the proper composite name
+  TODO: Wrap JIT in a `stablehlo.composite` op, instead of generating a call op.
+  """
+
+  def composite_impl(*args):
+    return jaxpr_impl(*args)
+
+  composite_impl.__name__ = composite_name
+  composite_impl.__qualname__ = composite_name
+  return jax.jit(composite_impl, **jit_args)
+
+
+def register_jax_composite(composite_name, impl, *ops, **jit_args):
+  """Register a composite using a JAX implementation.
+    composite_name - The name of the library op to use in the exported composite
+    impl           - A JAX lowering for the library operation
+    *ops           - Variadic torch.ops to lower using `impl`.
+    **jit_args     - Additional parameters to forward to JAX jit.
+
+  This is used to register custom lowerings with an explicit jaxpr
+  implementation, such as preserving a specific aten op using a jaten impl.
+
+  For custom torch op registration with a decomposition written in torch,
+  use `register_torch_composite`.
+
+  For jit params and troubleshooting see:
+  https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html
+  """
+
+  @jaten.op(*ops)
+  def _composite_impl(*args):
+    return _jit_composite_impl(composite_name, impl, **jit_args)(*args)
+
+
+def register_torch_composite(composite_name, impl, *ops, **jit_args):
+  """Register a torch decomposition as a composite.
+  This is useful for registerring custom torch op libraries as composite ops.
+
+  The `impl` can be the `@impl` used to define the torch custom library op.
+  This must be a function or module impl that provides the decompositions, and
+  not an instance of the custom op.
+
+  TODO: Better error handling, or can we make this an instance of the op as a param?
+
+  For jit params and troubleshooting see:
+  https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html
+  """
+
+  @jaten.op(*ops)
+  def _composite_impl(*args):
+
+    class ImplWrapper(torch.nn.Module):
+
+      def __init__(self):
+        super().__init__()
+
+      def forward(self, *args):
+        return impl(*args)
+
+    # Note: avoid refactoring to share code with register_jaxpr_composite.
+    # The `extract_jax` call must live in the `@jaten.op` handler. If called
+    # outside of the handler, we would build the jaxpr representation of the
+    # module once during registration, potentially missing op registrations that
+    # come after. I.e. may miss nested abstractions if we build jaxpr AoT.
+    state, jfn = torchax.extract_jax(ImplWrapper())
+    jaxpr_impl = lambda *args: jfn(state, tuple([*args]))
+    return _jit_composite_impl(composite_name, jaxpr_impl, **jit_args)(*args)