torchax 0.0.10.dev20251116__py3-none-any.whl → 0.0.11.dev202617__py3-none-any.whl

torchax/ops/jtorchvision_nms.py

@@ -16,12 +16,10 @@
 Forked at: https://raw.githubusercontent.com/mlperf/training_results_v0.7/refs/heads/master/Google/benchmarks/ssd/implementations/ssd-research-JAX-tpu-v3-4096/nms.py
 """
 
-import functools
-from typing import List, Union, Optional, Tuple
-
+import jax.numpy as jnp
 import torch
 from jax import lax
-import jax.numpy as jnp
+
 from . import ops_registry
 
 _NMS_TILE_SIZE = 256
@@ -38,9 +36,11 @@ def _bbox_overlap(boxes, gt_boxes):
     iou: Intersection over union matrix of all input bounding boxes
   """
   bb_y_min, bb_x_min, bb_y_max, bb_x_max = jnp.split(
-      ary=boxes, indices_or_sections=4, axis=2)
+    ary=boxes, indices_or_sections=4, axis=2
+  )
   gt_y_min, gt_x_min, gt_y_max, gt_x_max = jnp.split(
-      ary=gt_boxes, indices_or_sections=4, axis=2)
+    ary=gt_boxes, indices_or_sections=4, axis=2
+  )
 
   # Calculates the intersection area.
   i_xmin = jnp.maximum(bb_x_min, jnp.transpose(gt_x_min, [0, 2, 1]))
@@ -64,11 +64,16 @@ def _bbox_overlap(boxes, gt_boxes):
 def _self_suppression(in_args):
   iou, _, iou_sum = in_args
   batch_size = iou.shape[0]
-  can_suppress_others = jnp.reshape(
-      jnp.max(iou, 1) <= 0.5, [batch_size, -1, 1]).astype(iou.dtype)
-  iou_suppressed = jnp.reshape(
-      (jnp.max(can_suppress_others * iou, 1) <= 0.5).astype(
-          iou.dtype), [batch_size, -1, 1]) * iou
+  can_suppress_others = jnp.reshape(jnp.max(iou, 1) <= 0.5, [batch_size, -1, 1]).astype(
+    iou.dtype
+  )
+  iou_suppressed = (
+    jnp.reshape(
+      (jnp.max(can_suppress_others * iou, 1) <= 0.5).astype(iou.dtype),
+      [batch_size, -1, 1],
+    )
+    * iou
+  )
   iou_sum_new = jnp.sum(iou_suppressed, [1, 2])
   return iou_suppressed, jnp.any(iou_sum - iou_sum_new > 0.5), iou_sum_new
 
@@ -76,11 +81,14 @@ def _self_suppression(in_args):
 def _cross_suppression(in_args):
   boxes, box_slice, iou_threshold, inner_idx = in_args
   batch_size = boxes.shape[0]
-  new_slice = lax.dynamic_slice(boxes, [0, inner_idx * _NMS_TILE_SIZE, 0],
-                                [batch_size, _NMS_TILE_SIZE, 4])
+  new_slice = lax.dynamic_slice(
+    boxes, [0, inner_idx * _NMS_TILE_SIZE, 0], [batch_size, _NMS_TILE_SIZE, 4]
+  )
   iou = _bbox_overlap(new_slice, box_slice)
-  ret_slice = jnp.expand_dims((jnp.all(iou < iou_threshold, [1])).astype(
-      box_slice.dtype), 2) * box_slice
+  ret_slice = (
+    jnp.expand_dims((jnp.all(iou < iou_threshold, [1])).astype(box_slice.dtype), 2)
+    * box_slice
+  )
   return boxes, ret_slice, iou_threshold, inner_idx + 1
 
 
@@ -101,38 +109,44 @@ def _suppression_loop_body(in_args):
   batch_size = boxes.shape[0]
 
   # Iterates over tiles that can possibly suppress the current tile.
-  box_slice = lax.dynamic_slice(boxes, [0, idx * _NMS_TILE_SIZE, 0],
-                                [batch_size, _NMS_TILE_SIZE, 4])
+  box_slice = lax.dynamic_slice(
+    boxes, [0, idx * _NMS_TILE_SIZE, 0], [batch_size, _NMS_TILE_SIZE, 4]
+  )
 
   def _loop_cond(in_args):
     _, _, _, inner_idx = in_args
     return inner_idx < idx
 
-  _, box_slice, _, _ = lax.while_loop(_loop_cond, _cross_suppression,
-                                      (boxes, box_slice, iou_threshold, 0))
+  _, box_slice, _, _ = lax.while_loop(
+    _loop_cond, _cross_suppression, (boxes, box_slice, iou_threshold, 0)
+  )
 
   # Iterates over the current tile to compute self-suppression.
   iou = _bbox_overlap(box_slice, box_slice)
   mask = jnp.expand_dims(
-      jnp.reshape(jnp.arange(_NMS_TILE_SIZE), [1, -1])
-      > jnp.reshape(jnp.arange(_NMS_TILE_SIZE), [-1, 1]), 0)
+    jnp.reshape(jnp.arange(_NMS_TILE_SIZE), [1, -1])
+    > jnp.reshape(jnp.arange(_NMS_TILE_SIZE), [-1, 1]),
+    0,
+  )
   iou *= (jnp.logical_and(mask, iou >= iou_threshold)).astype(iou.dtype)
 
   def _loop_cond2(in_args):
     _, loop_condition, _ = in_args
     return loop_condition
 
-  suppressed_iou, _, _ = lax.while_loop(_loop_cond2, _self_suppression,
-                                        (iou, True, jnp.sum(iou, [1, 2])))
+  suppressed_iou, _, _ = lax.while_loop(
+    _loop_cond2, _self_suppression, (iou, True, jnp.sum(iou, [1, 2]))
+  )
   suppressed_box = jnp.sum(suppressed_iou, 1) > 0
   box_slice *= jnp.expand_dims(1.0 - suppressed_box.astype(box_slice.dtype), 2)
 
   # Uses box_slice to update the input boxes.
-  mask = jnp.reshape((jnp.equal(jnp.arange(num_tiles),
-                                idx)).astype(boxes.dtype), [1, -1, 1, 1])
-  boxes = jnp.tile(jnp.expand_dims(
-      box_slice, 1), [1, num_tiles, 1, 1]) * mask + jnp.reshape(
-          boxes, [batch_size, num_tiles, _NMS_TILE_SIZE, 4]) * (1 - mask)
+  mask = jnp.reshape(
+    (jnp.equal(jnp.arange(num_tiles), idx)).astype(boxes.dtype), [1, -1, 1, 1]
+  )
+  boxes = jnp.tile(
+    jnp.expand_dims(box_slice, 1), [1, num_tiles, 1, 1]
+  ) * mask + jnp.reshape(boxes, [batch_size, num_tiles, _NMS_TILE_SIZE, 4]) * (1 - mask)
   boxes = jnp.reshape(boxes, [batch_size, -1, 4])
 
   # Updates output_size.
@@ -193,8 +207,7 @@ def non_max_suppression_padded(scores, boxes, max_output_size, iou_threshold):
   """
   batch_size = boxes.shape[0]
   num_boxes = boxes.shape[1]
-  pad = int(jnp.ceil(
-      float(num_boxes) / _NMS_TILE_SIZE)) * _NMS_TILE_SIZE - num_boxes
+  pad = int(jnp.ceil(float(num_boxes) / _NMS_TILE_SIZE)) * _NMS_TILE_SIZE - num_boxes
   boxes = jnp.pad(boxes.astype(jnp.float32), [[0, 0], [0, pad], [0, 0]])
   scores = jnp.pad(scores.astype(jnp.float32), [[0, 0], [0, pad]])
   num_boxes += pad
@@ -202,29 +215,37 @@ def non_max_suppression_padded(scores, boxes, max_output_size, iou_threshold):
   def _loop_cond(in_args):
     unused_boxes, unused_threshold, output_size, idx = in_args
     return jnp.logical_and(
-        jnp.min(output_size) < max_output_size, idx
-        < num_boxes // _NMS_TILE_SIZE)
+      jnp.min(output_size) < max_output_size, idx < num_boxes // _NMS_TILE_SIZE
+    )
 
   selected_boxes, _, output_size, _ = lax.while_loop(
-      _loop_cond, _suppression_loop_body,
-      (boxes, iou_threshold, jnp.zeros([batch_size], jnp.int32), 0))
+    _loop_cond,
+    _suppression_loop_body,
+    (boxes, iou_threshold, jnp.zeros([batch_size], jnp.int32), 0),
+  )
   idx = num_boxes - lax.top_k(
-      jnp.any(selected_boxes > 0, [2]).astype(jnp.int32) *
-      jnp.expand_dims(jnp.arange(num_boxes, 0, -1), 0),
-      max_output_size)[0].astype(jnp.int32)
+    jnp.any(selected_boxes > 0, [2]).astype(jnp.int32)
+    * jnp.expand_dims(jnp.arange(num_boxes, 0, -1), 0),
+    max_output_size,
+  )[0].astype(jnp.int32)
   idx = jnp.minimum(idx, num_boxes - 1)
   idx = jnp.reshape(
-      idx + jnp.reshape(jnp.arange(batch_size) * num_boxes, [-1, 1]), [-1])
+    idx + jnp.reshape(jnp.arange(batch_size) * num_boxes, [-1, 1]), [-1]
+  )
 
   return idx
-  boxes = jnp.reshape((jnp.reshape(boxes, [-1, 4]))[idx],
-                      [batch_size, max_output_size, 4])
-  boxes = boxes * (jnp.reshape(jnp.arange(max_output_size), [1, -1, 1])
-                   < jnp.reshape(output_size, [-1, 1, 1])).astype(boxes.dtype)
-  scores = jnp.reshape(
-      jnp.reshape(scores, [-1, 1])[idx], [batch_size, max_output_size])
-  scores = scores * (jnp.reshape(jnp.arange(max_output_size), [1, -1])
-                     < jnp.reshape(output_size, [-1, 1])).astype(scores.dtype)
+  boxes = jnp.reshape(
+    (jnp.reshape(boxes, [-1, 4]))[idx], [batch_size, max_output_size, 4]
+  )
+  boxes = boxes * (
+    jnp.reshape(jnp.arange(max_output_size), [1, -1, 1])
+    < jnp.reshape(output_size, [-1, 1, 1])
+  ).astype(boxes.dtype)
+  scores = jnp.reshape(jnp.reshape(scores, [-1, 1])[idx], [batch_size, max_output_size])
+  scores = scores * (
+    jnp.reshape(jnp.arange(max_output_size), [1, -1])
+    < jnp.reshape(output_size, [-1, 1])
+  ).astype(scores.dtype)
   return scores, boxes
 
 
@@ -235,14 +256,13 @@ def nms(boxes, scores, iou_threshold):
   max_output_size = boxes.shape[0]
   boxes = boxes.reshape((1, *boxes.shape))
   scores = scores.reshape((1, *scores.shape))
-  res = non_max_suppression_padded(scores, boxes, max_output_size,
-                                   iou_threshold)
+  res = non_max_suppression_padded(scores, boxes, max_output_size, iou_threshold)
   return res
 
 
 try:
   import torch
-  import torchvision
+
   ops_registry.register_torch_dispatch_op(torch.ops.torchvision.nms, nms)
 except Exception:
   pass

torchax/ops/mappings.py

@@ -12,20 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from jax import dlpack as jaxdl
 import jax.numpy as jnp
 import numpy
 import torch
 import torch.func
-import torch.utils.dlpack as torchdl
 import torch.utils._mode_utils as mode_utils
+import torch.utils.dlpack as torchdl
+from jax import dlpack as jaxdl
 
 NUMPY_UNSUPPORTED_DTYPES = {
-    torch.bfloat16: jnp.bfloat16,
-    torch.float8_e4m3fn: jnp.float8_e4m3fn,
-    torch.float8_e4m3fnuz: jnp.float8_e4m3fnuz,
-    torch.float8_e5m2: jnp.float8_e5m2,
-    torch.float8_e5m2fnuz: jnp.float8_e5m2fnuz,
+  torch.bfloat16: jnp.bfloat16,
+  torch.float8_e4m3fn: jnp.float8_e4m3fn,
+  torch.float8_e4m3fnuz: jnp.float8_e4m3fnuz,
+  torch.float8_e5m2: jnp.float8_e5m2,
+  torch.float8_e5m2fnuz: jnp.float8_e5m2fnuz,
 }
 
 
@@ -51,8 +51,9 @@ def t2j(t, use_dlpack=True):
     # https://github.com/google/jax/issues/7657
     # https://github.com/google/jax/issues/17784
     if t.dtype in NUMPY_UNSUPPORTED_DTYPES:
-      nparray = (t.cpu().detach().to(torch.float32).numpy()
-                )  # handle dtypes not supported by numpy
+      nparray = (
+        t.cpu().detach().to(torch.float32).numpy()
+      )  # handle dtypes not supported by numpy
     else:
       nparray = t.cpu().detach().numpy()
     res = jnp.asarray(nparray)
@@ -91,71 +92,49 @@ def j2t(x, use_dlpack=True):
 
 
 TORCH_DTYPE_TO_JAX = {
-    # NO_MAPPING        : jnp.float0.dtype (signless scalar int),
-    torch.bool:
-        jnp.bool_.dtype,
-    # NO_MAPPING        : jnp.int4.dtype,
-    torch.int8:
-        jnp.int8.dtype,
-    torch.int16:
-        jnp.int16.dtype,
-    torch.int32:
-        jnp.int32.dtype,
-    torch.int64:
-        jnp.int64.dtype,
-    torch.long:
-        jnp.int64.dtype,
-    # NO_MAPPING        : jnp.uint4
-    torch.uint8:
-        jnp.uint8.dtype,
-    torch.uint16:
-        jnp.uint16.dtype,
-    torch.uint32:
-        jnp.uint32.dtype,
-    torch.uint64:
-        jnp.uint64.dtype,
-    # NO_MAPPING        : jnp.float8_e4m3b11fnuz.dtype,
-    torch.float8_e4m3fn:
-        jnp.float8_e4m3fn.dtype,
-    # NO_MAPPING        : jnp.float8_e4m3fnuz.dtype,
-    torch.float8_e5m2:
-        jnp.float8_e5m2.dtype,
-    # NO_MAPPING        : jnp.float8_e5m2fnuz.dtype,
-    torch.bfloat16:
-        jnp.bfloat16.dtype,
-    torch.half:
-        jnp.float16.dtype,
-    torch.float16:
-        jnp.float16.dtype,
-    torch.float32:
-        jnp.float32.dtype,
-    torch.float64:
-        jnp.float64.dtype,
-    torch.double:
-        jnp.double.dtype,
-    torch.complex64:
-        jnp.complex64.dtype,
-    torch.complex128:
-        jnp.complex128.dtype,
-    None:
-        None,
+  # NO_MAPPING        : jnp.float0.dtype (signless scalar int),
+  torch.bool: jnp.bool_.dtype,
+  # NO_MAPPING        : jnp.int4.dtype,
+  torch.int8: jnp.int8.dtype,
+  torch.int16: jnp.int16.dtype,
+  torch.int32: jnp.int32.dtype,
+  torch.int64: jnp.int64.dtype,
+  torch.long: jnp.int64.dtype,
+  # NO_MAPPING        : jnp.uint4
+  torch.uint8: jnp.uint8.dtype,
+  torch.uint16: jnp.uint16.dtype,
+  torch.uint32: jnp.uint32.dtype,
+  torch.uint64: jnp.uint64.dtype,
+  torch.float4_e2m1fn_x2: jnp.float4_e2m1fn.dtype,
+  # NO_MAPPING        : jnp.float8_e4m3b11fnuz.dtype,
+  torch.float8_e4m3fn: jnp.float8_e4m3fn.dtype,
+  # NO_MAPPING        : jnp.float8_e4m3fnuz.dtype,
+  torch.float8_e5m2: jnp.float8_e5m2.dtype,
+  # NO_MAPPING        : jnp.float8_e5m2fnuz.dtype,
+  torch.bfloat16: jnp.bfloat16.dtype,
+  torch.half: jnp.float16.dtype,
+  torch.float16: jnp.float16.dtype,
+  torch.float32: jnp.float32.dtype,
+  torch.float64: jnp.float64.dtype,
+  torch.double: jnp.double.dtype,
+  torch.complex64: jnp.complex64.dtype,
+  torch.complex128: jnp.complex128.dtype,
+  None: None,
 }
 
 JAX_DTYPE_TO_TORCH = {value: key for key, value in TORCH_DTYPE_TO_JAX.items()}
 # Add imprecise mappings for some JAX dtypes which don't have torch analogues
-JAX_DTYPE_TO_TORCH[jnp.dtype('int4')] = torch.int8
-JAX_DTYPE_TO_TORCH[jnp.dtype('uint4')] = torch.uint8
+JAX_DTYPE_TO_TORCH[jnp.dtype("int4")] = torch.int8
+JAX_DTYPE_TO_TORCH[jnp.dtype("uint4")] = torch.uint8
 
 
 def t2j_dtype(dtype):
   if dtype not in TORCH_DTYPE_TO_JAX:
-    raise RuntimeError(
-        f'Attempting to convert unknown type: {dtype} to jax type,')
+    raise RuntimeError(f"Attempting to convert unknown type: {dtype} to jax type,")
   return TORCH_DTYPE_TO_JAX[dtype]
 
 
 def j2t_dtype(dtype):
   if dtype not in JAX_DTYPE_TO_TORCH:
-    raise RuntimeError(
-        f'Attempting to convert unknown type: {dtype} to torch type,')
+    raise RuntimeError(f"Attempting to convert unknown type: {dtype} to torch type,")
   return JAX_DTYPE_TO_TORCH[dtype]

torchax/ops/op_base.py

@@ -13,25 +13,23 @@
 # limitations under the License.
 
 import functools
+from collections.abc import Callable
+from typing import Concatenate, ParamSpec
+
 import jax
 import jax.numpy as jnp
 import numpy as np
 import torch
+
+from torchax import types
 from torchax.ops import mappings
 from torchax.view import View
-from torchax import types
-import sys
-
-from typing import Callable, Optional, ParamSpec, Concatenate
 
 
 class InplaceOp:
-
-  def __init__(self,
-               functional_op,
-               replace=False,
-               position_to_mutate=0,
-               is_jax_func=False):
+  def __init__(
+    self, functional_op, replace=False, position_to_mutate=0, is_jax_func=False
+  ):
     self.functional = functional_op
     self.replace = replace
     self.position_to_mutate = position_to_mutate
@@ -65,15 +63,14 @@ class InplaceOp:
 
 
 class OutVariant:
-
   def __call__(self, *args, **kwargs):
-    to_mutate = kwargs['out']
-    del kwargs['out']
+    to_mutate = kwargs["out"]
+    del kwargs["out"]
     to_mutate._elem = self.functional(*args, **kwargs)._elem
     return to_mutate
 
 
-P = ParamSpec('P')
+P = ParamSpec("P")
 
 
 def convert_dtype(use_default_dtype: bool = True):
@@ -87,11 +84,8 @@ def convert_dtype(use_default_dtype: bool = True):
   """
 
   def decorator(func: types.TorchCallable):
-
     @functools.wraps(func)
-    def wrapper(*args: P.args,
-                dtype: Optional[torch.dtype] = None,
-                **kwargs: P.kwargs):
+    def wrapper(*args: P.args, dtype: torch.dtype | None = None, **kwargs: P.kwargs):
       if not dtype and use_default_dtype:
         dtype = torch.get_default_dtype()
       if isinstance(dtype, torch.dtype):
@@ -106,8 +100,7 @@ def convert_dtype(use_default_dtype: bool = True):
   return decorator
 
 
-def maybe_convert_constant_dtype(val: Optional[types.JaxValue],
-                                 dtype: Optional[jnp.dtype]):
+def maybe_convert_constant_dtype(val: types.JaxValue | None, dtype: jnp.dtype | None):
   """Optionally converts scalar constant's dtype using `numpy`
 
   Use in cases where you require a constant and can't handle a traced array.
@@ -134,12 +127,11 @@ def promote_int_input(f: Callable[Concatenate[jax.Array, P], types.JaxValue]):
   return wrapper
 
 
-def foreach_loop(seq: jax.Array,
-                 fn: Callable[[jax.Array, jax.Array], jax.Array],
-                 init_val=0.0):
+def foreach_loop(
+  seq: jax.Array, fn: Callable[[jax.Array, jax.Array], jax.Array], init_val=0.0
+):
   """Run `fn` for each element of 1D array `seq`.
 
   Similar to `functools.reduce`, but implemented with `jax.lax.fori_loop`."""
   assert len(seq.shape) == 1
-  return jax.lax.fori_loop(0, len(seq), lambda i, carry: fn(carry, seq[i]),
-                           init_val)
+  return jax.lax.fori_loop(0, len(seq), lambda i, carry: fn(carry, seq[i]), init_val)

torchax/ops/ops_registry.py

@@ -14,56 +14,61 @@
 
 import dataclasses
 import logging
-from torchax.types import JaxCallable, TorchCallable
 
-from typing import Union, Dict
+from torchax.types import JaxCallable, TorchCallable
 
 
 @dataclasses.dataclass
 class Operator:
   torch_op: TorchCallable
-  func: Union[TorchCallable, JaxCallable]
+  func: TorchCallable | JaxCallable
   is_jax_function: bool
   is_user_defined: bool
   needs_env: bool
   is_view_op: bool
 
 
-all_aten_ops: Dict[TorchCallable, Operator] = {}
-all_torch_functions: Dict[TorchCallable, Operator] = {}
+all_aten_ops: dict[TorchCallable, Operator] = {}
+all_torch_functions: dict[TorchCallable, Operator] = {}
 
 
-def register_torch_dispatch_op(aten_op,
-                               impl_callable,
-                               is_jax_function=True,
-                               is_user_defined=False,
-                               needs_env=False,
-                               is_view_op=False):
+def register_torch_dispatch_op(
+  aten_op,
+  impl_callable,
+  is_jax_function=True,
+  is_user_defined=False,
+  needs_env=False,
+  is_view_op=False,
+):
   op = Operator(
-      aten_op,
-      impl_callable,
-      is_jax_function=is_jax_function,
-      is_user_defined=is_user_defined,
-      needs_env=needs_env,
-      is_view_op=is_view_op)
+    aten_op,
+    impl_callable,
+    is_jax_function=is_jax_function,
+    is_user_defined=is_user_defined,
+    needs_env=needs_env,
+    is_view_op=is_view_op,
+  )
   if aten_op in all_aten_ops:
-    logging.warning(f'Duplicate op registration for {aten_op}')
+    logging.warning(f"Duplicate op registration for {aten_op}")
   all_aten_ops[aten_op] = op
   return impl_callable
 
 
-def register_torch_function_op(torch_func,
-                               impl_callable,
-                               is_jax_function=True,
-                               is_user_defined=False,
-                               needs_env=False,
-                               is_view_op=False):
+def register_torch_function_op(
+  torch_func,
+  impl_callable,
+  is_jax_function=True,
+  is_user_defined=False,
+  needs_env=False,
+  is_view_op=False,
+):
   op = Operator(
-      torch_func,
-      impl_callable,
-      is_jax_function=is_jax_function,
-      is_user_defined=is_user_defined,
-      needs_env=needs_env,
-      is_view_op=is_view_op)
+    torch_func,
+    impl_callable,
+    is_jax_function=is_jax_function,
+    is_user_defined=is_user_defined,
+    needs_env=needs_env,
+    is_view_op=is_view_op,
+  )
   all_torch_functions[torch_func] = op
   return impl_callable