torchax 0.0.10.dev20251118__py3-none-any.whl

torchax/ops/jax_reimplement.py

@@ -0,0 +1,211 @@
+# 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 Callable, Sequence
+
+import numpy as np
+from jax import lax
+from jax import numpy as jnp
+from jax._src import core
+from jax._src.image.scale import ResizeMethod, _kernels
+from jax._src.numpy.util import promote_dtypes_inexact
+from jax._src.util import canonicalize_axis
+
+# 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.0 / 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.0) if antialias else 1.0
+  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.0 * 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,64 @@
+# 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
+import torch
+
+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,122 @@
+# 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 jax
+import torch
+
+import torchax
+from torchax.ops import jaten
+
+
+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())
+
+    def jaxpr_impl(*args):
+      return jfn(state, (*args,))
+
+    return _jit_composite_impl(composite_name, jaxpr_impl, **jit_args)(*args)