torchax 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

torchax/ops/jax_reimplement.py

@@ -1,4 +1,3 @@
-
 from collections.abc import Sequence
 from jax._src.numpy.util import promote_dtypes_inexact
 import numpy as np
@@ -15,12 +14,9 @@ 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
   # When downsampling the kernel should be scaled since we want to low pass
@@ -38,8 +34,8 @@ def compute_weight_mat(input_size: core.DimSize,
   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.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.
@@ -48,12 +44,14 @@ def compute_weight_mat(input_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):
@@ -70,8 +68,8 @@ 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
@@ -81,15 +79,19 @@ def _scale_and_translate(x, output_shape: core.Shape,
 
 # 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):
+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
@@ -165,5 +167,5 @@ def scale_and_translate(image, shape: core.Shape,
   return _scale_and_translate(image, shape, spatial_dims, scale, translation,
                               kernel, antialias, precision)
 
-# END ----------------- END JAX code copied for testing -----------------------------
 
+# END ----------------- END JAX code copied for testing -----------------------------

torchax/ops/jc10d.py

@@ -6,6 +6,7 @@ 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)
@@ -21,7 +22,7 @@ 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":
@@ -38,9 +39,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

@@ -0,0 +1,113 @@
+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

@@ -14,19 +14,22 @@ 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.
 
@@ -36,10 +39,12 @@ def register_jax_composite(composite_name, impl, *ops, **jit_args):
   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.
@@ -53,10 +58,12 @@ def register_torch_composite(composite_name, impl, *ops, **jit_args):
   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__()