torchax 0.0.10.dev20251118__py3-none-any.whl

torchax/ops/jtorchvision_nms.py

@@ -0,0 +1,268 @@
+# 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.
+
+"""
+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 jax.numpy as jnp
+import torch
+from jax import lax
+
+from . import ops_registry
+
+_NMS_TILE_SIZE = 256
+
+
+def _bbox_overlap(boxes, gt_boxes):
+  """Find Bounding box overlap.
+
+  Args:
+    boxes: first set of bounding boxes
+    gt_boxes: second set of boxes to compute IOU
+
+  Returns:
+    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
+  )
+  gt_y_min, gt_x_min, gt_y_max, gt_x_max = jnp.split(
+    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]))
+  i_xmax = jnp.minimum(bb_x_max, jnp.transpose(gt_x_max, [0, 2, 1]))
+  i_ymin = jnp.maximum(bb_y_min, jnp.transpose(gt_y_min, [0, 2, 1]))
+  i_ymax = jnp.minimum(bb_y_max, jnp.transpose(gt_y_max, [0, 2, 1]))
+  i_area = jnp.maximum((i_xmax - i_xmin), 0) * jnp.maximum((i_ymax - i_ymin), 0)
+
+  # Calculates the union area.
+  bb_area = (bb_y_max - bb_y_min) * (bb_x_max - bb_x_min)
+  gt_area = (gt_y_max - gt_y_min) * (gt_x_max - gt_x_min)
+  # Adds a small epsilon to avoid divide-by-zero.
+  u_area = bb_area + jnp.transpose(gt_area, [0, 2, 1]) - i_area + 1e-8
+
+  # Calculates IoU.
+  iou = i_area / u_area
+
+  return iou
+
+
+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
+  )
+  iou_sum_new = jnp.sum(iou_suppressed, [1, 2])
+  return iou_suppressed, jnp.any(iou_sum - iou_sum_new > 0.5), iou_sum_new
+
+
+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]
+  )
+  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
+  )
+  return boxes, ret_slice, iou_threshold, inner_idx + 1
+
+
+def _suppression_loop_body(in_args):
+  """Process boxes in the range [idx*_NMS_TILE_SIZE, (idx+1)*_NMS_TILE_SIZE).
+
+  Args:
+    in_args: A tuple of arguments: boxes, iou_threshold, output_size, idx
+
+  Returns:
+    boxes: updated boxes.
+    iou_threshold: pass down iou_threshold to the next iteration.
+    output_size: the updated output_size.
+    idx: the updated induction variable.
+  """
+  boxes, iou_threshold, output_size, idx = in_args
+  num_tiles = boxes.shape[1] // _NMS_TILE_SIZE
+  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]
+  )
+
+  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)
+  )
+
+  # 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,
+  )
+  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_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)
+  boxes = jnp.reshape(boxes, [batch_size, -1, 4])
+
+  # Updates output_size.
+  output_size += jnp.sum(jnp.any(box_slice > 0, [2]).astype(jnp.int32), [1])
+  return boxes, iou_threshold, output_size, idx + 1
+
+
+def non_max_suppression_padded(scores, boxes, max_output_size, iou_threshold):
+  """A wrapper that handles non-maximum suppression.
+
+  Assumption:
+    * The boxes are sorted by scores unless the box is a dot (all coordinates
+      are zero).
+    * Boxes with higher scores can be used to suppress boxes with lower scores.
+
+  The overal design of the algorithm is to handle boxes tile-by-tile:
+
+  boxes = boxes.pad_to_multiply_of(tile_size)
+  num_tiles = len(boxes) // tile_size
+  output_boxes = []
+  for i in range(num_tiles):
+    box_tile = boxes[i*tile_size : (i+1)*tile_size]
+    for j in range(i - 1):
+      suppressing_tile = boxes[j*tile_size : (j+1)*tile_size]
+      iou = _bbox_overlap(box_tile, suppressing_tile)
+      # if the box is suppressed in iou, clear it to a dot
+      box_tile *= _update_boxes(iou)
+    # Iteratively handle the diagnal tile.
+    iou = _box_overlap(box_tile, box_tile)
+    iou_changed = True
+    while iou_changed:
+      # boxes that are not suppressed by anything else
+      suppressing_boxes = _get_suppressing_boxes(iou)
+      # boxes that are suppressed by suppressing_boxes
+      suppressed_boxes = _get_suppressed_boxes(iou, suppressing_boxes)
+      # clear iou to 0 for boxes that are suppressed, as they cannot be used
+      # to suppress other boxes any more
+      new_iou = _clear_iou(iou, suppressed_boxes)
+      iou_changed = (new_iou != iou)
+      iou = new_iou
+    # remaining boxes that can still suppress others, are selected boxes.
+    output_boxes.append(_get_suppressing_boxes(iou))
+    if len(output_boxes) >= max_output_size:
+      break
+
+  Args:
+    scores: a tensor with a shape of [batch_size, anchors].
+    boxes: a tensor with a shape of [batch_size, anchors, 4].
+    max_output_size: a scalar integer `Tensor` representing the maximum number
+      of boxes to be selected by non max suppression.
+    iou_threshold: a float representing the threshold for deciding whether boxes
+      overlap too much with respect to IOU.
+  Returns:
+    nms_scores: a tensor with a shape of [batch_size, anchors]. It has same
+      dtype as input scores.
+    nms_proposals: a tensor with a shape of [batch_size, anchors, 4]. It has
+      same dtype as input boxes.
+  """
+  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
+  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
+
+  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
+    )
+
+  selected_boxes, _, output_size, _ = lax.while_loop(
+    _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)
+  idx = jnp.minimum(idx, num_boxes - 1)
+  idx = jnp.reshape(
+    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)
+  return scores, boxes
+
+
+# registry:
+
+
+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)
+  return res
+
+
+try:
+  import torch
+
+  ops_registry.register_torch_dispatch_op(torch.ops.torchvision.nms, nms)
+except Exception:
+  pass

torchax/ops/mappings.py

@@ -0,0 +1,139 @@
+# 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.numpy as jnp
+import numpy
+import torch
+import torch.func
+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,
+}
+
+
+def t2j(t, use_dlpack=True):
+  is_bool = False
+  if t.dtype == torch.bool:
+    is_bool = True
+    t = t.to(torch.int8)
+
+  t = t.to_dense()
+
+  if not t.is_contiguous():
+    t = t.contiguous()
+
+  res = None
+  if use_dlpack:
+    try:
+      res = jaxdl.from_dlpack(t)
+    except Exception:
+      pass
+
+  if res is None:
+    # 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
+    else:
+      nparray = t.cpu().detach().numpy()
+    res = jnp.asarray(nparray)
+    if t.dtype in NUMPY_UNSUPPORTED_DTYPES:
+      res = res.astype(NUMPY_UNSUPPORTED_DTYPES[t.dtype])
+
+  if is_bool:
+    res = res.astype(jnp.bool_)
+  return res
+
+
+def j2t(x, use_dlpack=True):
+  with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
+    res = None
+    if use_dlpack:
+      try:
+        dl = jaxdl.to_dlpack(x)
+        res = torchdl.from_dlpack(dl)
+      except Exception:
+        res = None
+
+    orig_dtype = None
+    if res is None:
+      orig_dtype = None
+      if x.dtype == jnp.bfloat16.dtype:
+        orig_dtype = x.dtype
+        x = x.astype(jnp.float32.dtype)
+      res = torch.from_numpy(numpy.asarray(x))
+
+    if x.dtype == jnp.bool_:
+      res = res.to(torch.bool)
+
+    if orig_dtype is not None:
+      res = res.to(j2t_dtype(orig_dtype))
+    return res
+
+
+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,
+}
+
+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
+
+
+def t2j_dtype(dtype):
+  if dtype not in TORCH_DTYPE_TO_JAX:
+    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,")
+  return JAX_DTYPE_TO_TORCH[dtype]

torchax/ops/op_base.py

@@ -0,0 +1,137 @@
+# 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 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
+
+
+class InplaceOp:
+  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
+    self.is_jax_func = is_jax_func
+
+  def __call__(self, *args, **kwargs):
+    to_mutate = args[self.position_to_mutate]
+    view_value = to_mutate
+    if isinstance(to_mutate, View):
+      view_value = to_mutate.torch()
+      # Convert the target View to a Tensor, and
+      # leave the rest args as is. If other args are
+      # also View, they will be converted to tensors
+      # in the self.functional dispatch.
+    env = view_value._env
+    if self.is_jax_func:
+      view_value, args, kwargs = env.t2j_iso((view_value, args, kwargs))
+      new_value_jax = self.functional(view_value, *args[1:], **kwargs)
+      new_value = env.j2t_iso(new_value_jax)
+    else:
+      new_value = self.functional(view_value, *args[1:], **kwargs)
+
+    if isinstance(to_mutate, View):
+      to_mutate.update(new_value)
+    else:
+      if self.replace:
+        to_mutate._elem = new_value._elem
+      else:
+        to_mutate.copy_(new_value)
+    return to_mutate
+
+
+class OutVariant:
+  def __call__(self, *args, **kwargs):
+    to_mutate = kwargs["out"]
+    del kwargs["out"]
+    to_mutate._elem = self.functional(*args, **kwargs)._elem
+    return to_mutate
+
+
+P = ParamSpec("P")
+
+
+def convert_dtype(use_default_dtype: bool = True):
+  """Converts `dtype` kwarg of function from torch to JAX.
+
+  Args:
+    use_default_dtype: Whether to use torch default dtype if none is provided.
+
+  Returns:
+    A decorator that wraps a JAX implementation of a torch function.
+  """
+
+  def decorator(func: types.TorchCallable):
+    @functools.wraps(func)
+    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):
+        jax_dtype = mappings.t2j_dtype(dtype)
+      else:
+        jax_dtype = dtype
+
+      return func(*args, dtype=jax_dtype, **kwargs)
+
+    return wrapper
+
+  return decorator
+
+
+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.
+  """
+  if val and dtype:
+    if isinstance(val, jax.Array):
+      return maybe_convert_constant_dtype(val.item(), dtype)
+
+    return np.array(val, dtype)
+
+  return val
+
+
+def promote_int_input(f: Callable[Concatenate[jax.Array, P], types.JaxValue]):
+  """If the first argument is an int array, promote it to float32."""
+
+  @functools.wraps(f)
+  def wrapper(x: jax.Array, *args: P.args, **kwargs: P.kwargs):
+    if x.dtype in [jnp.int8, jnp.int16, jnp.int32, jnp.int64]:
+      x = x.astype(mappings.t2j_dtype(torch.get_default_dtype()))
+
+    return f(x, *args, **kwargs)
+
+  return wrapper
+
+
+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)

torchax/ops/ops_registry.py

@@ -0,0 +1,74 @@
+# 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 dataclasses
+import logging
+
+from torchax.types import JaxCallable, TorchCallable
+
+
+@dataclasses.dataclass
+class Operator:
+  torch_op: TorchCallable
+  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] = {}
+
+
+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,
+  )
+  if aten_op in all_aten_ops:
+    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,
+):
+  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,
+  )
+  all_torch_functions[torch_func] = op
+  return impl_callable