torchax 0.0.10.dev20251114__py3-none-any.whl → 0.0.11.dev202612__py3-none-any.whl

torchax/ops/jax_reimplement.py

@@ -12,16 +12,15 @@
 # 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
+from collections.abc import Callable, Sequence
+
 import numpy as np
-import jax
+from jax import lax
 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
+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 -----------------------------
@@ -29,27 +28,39 @@ from typing import Callable
 # 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):
+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
+  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.) if antialias else 1.
-  sample_f = ((jnp.arange(output_size, dtype=dtype) + 0.5) * inv_scale -
-              translation * inv_scale - 0.5)
+  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)
+    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)
+    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.
@@ -58,17 +69,26 @@ def compute_weight_mat(input_size: core.DimSize, output_size: core.DimSize,
   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)
+    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):
+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)
@@ -82,8 +102,9 @@ def _scale_and_translate(x, output_shape: core.Shape,
     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)
+    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
@@ -97,15 +118,16 @@ def _scale_and_translate(x, output_shape: core.Shape,
 # 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):
+  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
@@ -163,23 +185,27 @@ def scale_and_translate(
   """
   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}')
+    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.')
+    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)
+  (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)
+  return _scale_and_translate(
+    image, shape, spatial_dims, scale, translation, kernel, antialias, precision
+  )
 
 
 # END ----------------- END JAX code copied for testing -----------------------------

torchax/ops/jc10d.py

@@ -12,15 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
 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)
@@ -36,7 +35,6 @@ def _c10d_all_gather(input, group_size: int, group_name: str):
 
 @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":
@@ -53,9 +51,9 @@ def _c10d_all_reduce(self, reduceOp: str, group_name: str):
 @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),
+    jax.lax.axis_index("torch_dist") == src,
+    self,
+    jnp.zeros_like(self),
   )
   return jax.lax.psum(masked, "torch_dist")
 

torchax/ops/jimage.py

@@ -21,19 +21,16 @@ def cubic_kernel(x, a=-0.75):
   absx = jnp.abs(x)
   x2 = absx * absx
   x3 = x2 * absx
-  cond1 = (absx <= 1)
+  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):
+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)
@@ -62,10 +59,10 @@ def gather_weights(img, idxs, 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
-    """
+  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
@@ -80,13 +77,11 @@ def interpolate_along_axis_bchw(img, idxs, weights, axis):
     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)
+    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)
+  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
@@ -108,20 +103,20 @@ def interpolate_bicubic_no_aa(img, out_h, out_w, align_corners=False):
     scale_x = out_w / w
 
   idxs_y, weights_y = compute_contribs(
-      h,
-      out_h,
-      scale_y,
-      align_corners=align_corners,
-      dtype=img.dtype,
+    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,
+    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

@@ -16,12 +16,11 @@
 during export. This includes aten ops, and custom operations.
 """
 
+import jax
 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):
@@ -75,9 +74,7 @@ def register_torch_composite(composite_name, impl, *ops, **jit_args):
 
   @jaten.op(*ops)
   def _composite_impl(*args):
-
     class ImplWrapper(torch.nn.Module):
-
       def __init__(self):
         super().__init__()
 
@@ -90,5 +87,8 @@ def register_torch_composite(composite_name, impl, *ops, **jit_args):
     # 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]))
+
+    def jaxpr_impl(*args):
+      return jfn(state, (*args,))
+
     return _jit_composite_impl(composite_name, jaxpr_impl, **jit_args)(*args)