torchax 0.0.4__py3-none-any.whl

torchax/ops/jtorch.py

@@ -0,0 +1,427 @@
+"""Tensor constructor overrides"""
+import math
+import collections.abc
+import functools
+from typing import Optional, Sequence
+import numpy as np
+
+import jax
+import jax.numpy as jnp
+from jax.experimental.pallas.ops.tpu import flash_attention
+from jax.experimental.shard_map import shard_map
+
+import torch
+from torchax.ops.ops_registry import register_torch_function_op
+from torchax.ops import op_base, mappings, jaten
+import torchax.tensor
+
+
+def register_function(torch_func, **kwargs):
+  return functools.partial(register_torch_function_op, torch_func, **kwargs)
+
+
+@register_function(torch.as_tensor, is_jax_function=False, needs_env=True)
+@op_base.convert_dtype(use_default_dtype=False)  # Attempt to infer type from elements
+def _as_tensor(data, dtype=None, device=None, env=None):
+  if isinstance(data, torch.Tensor):
+    return env._to_copy(data, dtype, device)
+  if isinstance(data, np.ndarray):
+    jax_res = jnp.asarray(data)
+  else:
+    jax_res = _tensor(data, dtype=dtype)
+  return torchax.tensor.Tensor(jax_res, env)
+
+
+@register_function(torch.tensor)
+@op_base.convert_dtype(use_default_dtype=False)  # Attempt to infer type from elements
+def _tensor(data, *, dtype=None, **kwargs):
+  python_types_to_torch_types = {
+      bool: jnp.bool,
+      int: jnp.int64,
+      float: jnp.float32,
+      complex: jnp.complex64,
+  }
+  if not dtype:
+    leaves = jax.tree_util.tree_leaves(data)
+    if len(leaves) > 0:
+      dtype = python_types_to_torch_types.get(type(leaves[0]))
+
+  return jnp.array(
+      data, dtype=dtype or mappings.t2j_dtype(torch.get_default_dtype()))
+
+
+@register_function(torch.allclose)
+def _aten_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
+  return jnp.allclose(input, other, rtol, atol, equal_nan)
+
+
+@register_function(torch.angle)
+def _torch_angle(input):
+  if input.dtype.name == 'int64':
+    input = input.astype(jnp.dtype('float32'))
+  return jnp.angle(input)
+
+
+@register_function(torch.argsort)
+def _torch_argsort(input, dim=-1, descending=False, stable=False):
+  expanded = False
+  if input.ndim == 0:
+    # for self of rank 0:
+    # torch.any(x, 0), torch.any(x, -1) works;
+    # torch.any(x, 1) throws out of bounds, so it's
+    # behavior is the same as a jnp array of rank 1
+    expanded = True
+    input = jnp.expand_dims(input, 0)
+  res = jnp.argsort(input, axis=dim, descending=descending,
+                     stable=stable)
+  if expanded:
+    res = res.squeeze()
+  return res
+
+@register_function(torch.diag)
+def _diag(input, diagonal=0):
+  return jnp.diag(input, k=diagonal)
+
+@register_function(torch.einsum)
+@register_function(torch.ops.aten.einsum)
+def _einsum(equation, *operands):
+  def get_params(*a):
+    inner_list = a[0]
+    if not isinstance(inner_list, jax.Array):
+      if len(inner_list) == 1:
+        A = inner_list
+        return A
+      elif len(inner_list) == 2:
+        A, B = inner_list
+        return A, B
+    return operands
+  assert isinstance(equation, str), 'Only accept str equation'
+  filtered_operands = get_params(*operands)
+  return jnp.einsum(equation, *filtered_operands)
+
+
+def _sdpa_reference(query, key, value, attn_mask=None, dropout_p=0.0,
+        is_causal=False, scale=None, enable_gqa=False) -> torch.Tensor:
+    L, S = query.size(-2), key.size(-2)
+    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
+    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
+    if is_causal:
+        assert attn_mask is None
+        temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=0)
+        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+        attn_bias.to(query.dtype)
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+        else:
+            attn_bias += attn_mask
+    if enable_gqa:
+        key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
+        value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
+
+    attn_weight = query @ key.transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    if dropout_p > 0:
+      attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+    return attn_weight @ value
+
+
+from jax.sharding import PartitionSpec
+
+def _tpu_flash_attention(query, key, value, env):
+  fsdp_partition = PartitionSpec('fsdp')
+  def wrap_flash_attention(query, key, value):
+    block_sizes = flash_attention.BlockSizes(
+      block_b=min(2, query.shape[0]),
+      block_q=min(512, query.shape[2]),
+      block_k_major=min(512, key.shape[2]),
+      block_k=min(512, key.shape[2]),
+      block_q_major_dkv=min(512, query.shape[2]),
+      block_k_major_dkv=min(512, key.shape[2]),
+      block_k_dkv=min(512, key.shape[2]),
+      block_q_dkv=min(512, query.shape[2]),
+      block_k_major_dq=min(512, key.shape[2]),
+      block_k_dq=min(256, key.shape[2]),
+      block_q_dq=min(1024, query.shape[2]),
+    )
+    return flash_attention.flash_attention(
+        query, key, value, causal=True, block_sizes=block_sizes)
+
+  if env.config.shmap_flash_attention:
+    wrap_flash_attention = shard_map(
+      wrap_flash_attention,
+      mesh=env._mesh,
+      in_specs=(fsdp_partition, fsdp_partition, fsdp_partition),
+      out_specs=fsdp_partition ,
+      check_rep=False,
+    )
+  #return flash_attn_mapped(query, key, value)
+  return wrap_flash_attention(query, key, value)
+
+
+@register_function(torch.nn.functional.pad)
+def pad(tensor, pad, mode="constant", value=None):
+  # For padding modes that have different names between Torch and NumPy, this
+  # dict provides a Torch-to-NumPy translation. Any string not in this dict will
+  # be passed through as-is.
+  MODE_NAME_TRANSLATION = {
+      "circular": "wrap",
+      "replicate": "edge",
+  }
+
+  numpy_mode = MODE_NAME_TRANSLATION.get(mode, mode)
+
+  num_prefix_dims = tensor.ndim - len(pad) // 2
+
+  numpy_pad_width = [(0, 0)] * num_prefix_dims
+  nd_slice = [slice(None)] * num_prefix_dims
+
+  for i in range(len(pad) - 2, -1, -2):
+    pad_start, pad_end = pad[i:i + 2]
+    slice_start, slice_end = None, None
+
+    if pad_start < 0:
+      slice_start = -pad_start
+      pad_start = 0
+
+    if pad_end < 0:
+      slice_end = pad_end
+      pad_end = 0
+
+    numpy_pad_width.append((pad_start, pad_end))
+    nd_slice.append(slice(slice_start, slice_end))
+
+  nd_slice = tuple(nd_slice)
+
+  # `jax.numpy.pad` complains if we provide an irrelevant `constant_values` arg,
+  # even if the value we pass in is `None`. (It treats `None` as `nan`.)
+  kwargs = dict()
+  if mode == "constant" and value is not None:
+    kwargs["constant_values"] = value
+
+  # The "replicate" mode pads first and then slices, whereas the "circular" mode
+  # slices first and then pads. The latter approach deals with smaller tensors,
+  # so we default to that option in modes where the order of operations doesn't
+  # affect the result.
+  if mode == "replicate":
+    return jnp.pad(tensor, numpy_pad_width, mode=numpy_mode, **kwargs)[nd_slice]
+  else:
+    return jnp.pad(tensor[nd_slice], numpy_pad_width, mode=numpy_mode, **kwargs)
+
+
+@register_function(torch.nn.functional.scaled_dot_product_attention, is_jax_function=False, needs_env=True)
+@register_function(torch.ops.aten.scaled_dot_product_attention, is_jax_function=False, needs_env=True)
+def scaled_dot_product_attention(
+   query, key, value, attn_mask=None,
+   dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, env=None) -> torch.Tensor:
+
+   if env.config.use_tpu_flash_attention:
+    jquery, jkey, jvalue = env.t2j_iso((query, key, value))
+    res = _tpu_flash_attention(jquery, jkey, jvalue, env)
+    return env.j2t_iso(res)
+
+   return _sdpa_reference(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa)
+
+@register_function(torch.Tensor.__getitem__)
+def getitem(self, indexes):
+  if isinstance(indexes, list) and isinstance(indexes[0], int):
+    # list of int, i.e. x[[1, 2]] NOT x[1, 2] (the second would be tuple of int)
+    indexes = (indexes, )
+  elif isinstance(indexes, list):
+    indexes = tuple(indexes)
+  return self[indexes]
+
+@register_function(torch.corrcoef)
+def _corrcoef(x):
+  if x.dtype.name == "int64":
+    return jnp.corrcoef(x).astype(jnp.float32)
+  return jnp.corrcoef(x)
+
+@register_function(torch.sparse.mm, is_jax_function=False)
+def _sparse_mm(mat1, mat2, reduce='sum'):
+  return torch.mm(mat1, mat2)
+
+@register_function(torch.isclose)
+def _aten_isclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
+  return jnp.isclose(input, other, rtol, atol, equal_nan)
+
+
+@register_function(torch.ones)
+def _ones(*size: int, dtype=None, **kwargs):
+  if len(size) == 1 and isinstance(size[0], collections.abc.Iterable):
+    size = size[0]
+  return jaten._ones(size, dtype=dtype)
+
+
+@register_function(torch.zeros, is_jax_function=True)
+def _zeros(*size: int, dtype=None, **kwargs):
+  if len(size) == 1 and isinstance(size[0], collections.abc.Iterable):
+    size = size[0]
+  return jaten._zeros(size, dtype=dtype)
+
+
+@register_function(torch.eye)
+@op_base.convert_dtype()
+def _eye(n: int, m: Optional[int] = None, *, dtype=None, **kwargs):
+  return jnp.eye(n, m, dtype=dtype)
+
+
+@register_function(torch.full)
+@op_base.convert_dtype()
+def _full(size: Sequence[int], fill_value, *, dtype=None, **kwargs):
+  # TODO: handle torch.Size
+  return jnp.full(size, fill_value, dtype=dtype)
+
+
+@register_function(torch.empty)
+@op_base.convert_dtype()
+def empty(*size: Sequence[int], dtype=None, **kwargs):
+  if len(size) == 1 and isinstance(size[0], collections.abc.Iterable):
+    size = size[0]
+  return jnp.empty(size, dtype=dtype)
+
+@register_function(torch.arange, is_jax_function=False)
+def arange(
+  start, end=None, step=None, 
+  out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False,
+  pin_memory=None,
+):
+  if end is None:
+    end = start
+    start = 0
+  if step is None:
+    step = 1
+  return torch.ops.aten.arange(start, end, step, dtype=dtype)
+
+@register_function(torch.empty_strided, is_jax_function=False)
+def empty_strided(
+  size, stride, *, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False):
+  return empty(size, dtype=dtype)
+
+
+@register_function(torch.unravel_index)
+def unravel_index(indices, shape):
+  return jnp.unravel_index(indices, shape)
+
+
+@register_function(torch.rand, is_jax_function=False)
+def rand(
+  *size, **kwargs
+):
+  if len(size) == 1 and isinstance(size[0], collections.abc.Iterable):
+    size = size[0]
+  return torch.ops.aten.rand(size, **kwargs)
+
+@register_function(torch.randn, is_jax_function=False)
+def randn(
+  *size, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False
+):
+  if len(size) == 1 and isinstance(size[0], collections.abc.Iterable):
+    size = size[0]
+  return torch.ops.aten.randn(size, generator=generator, dtype=dtype)
+
+@register_function(torch.randint, is_jax_function=False)
+def randint(
+  *args, **kwargs
+):
+  return torch.ops.aten.randint(*args, **kwargs)
+
+
+@register_function(torch.logdet)
+def logdet(input):
+  _, logabsdet = jaten._aten__linalg_slogdet(input)
+  return logabsdet
+
+
+@register_function(torch.linalg.slogdet)
+def linalg_slogdet(input):
+  sign, logabsdet = jaten._aten__linalg_slogdet(input)
+  return torch.return_types.slogdet((sign, logabsdet))
+
+
+@register_function(torch.tensor_split)
+def tensor_split(input, indices_or_sections, dim=0):
+  return jnp.array_split(input, indices_or_sections, axis=dim)
+
+
+@register_function(torch.linalg.solve)
+def linalg_solve(a, b):
+  res, _ = jaten._aten__linalg_solve_ex(a, b)
+  return res
+
+
+@register_function(torch.linalg.solve_ex)
+def linalg_solve_ex(a, b):
+  res, info = jaten._aten__linalg_solve_ex(a, b)
+  return res, info
+
+@register_function(torch.linalg.svd)
+def linalg_svd(a, full_matrices=True):
+  return jaten._aten__linalg_svd(a, full_matrices=full_matrices)
+
+@register_function(torch.linalg.matrix_power)
+def matrix_power(A, n, *, out=None):
+  return jnp.linalg.matrix_power(A, n)
+
+@register_function(torch.svd)
+def svd(a, some=True, compute_uv=True):
+  if not compute_uv:
+    S = jaten._aten__linalg_svd(a, full_matrices=False)[1]
+    U = jnp.zeros((a.shape[-2], a.shape[-2]), dtype=a.dtype)
+    V = jnp.zeros((a.shape[-1], a.shape[-1]), dtype=a.dtype)
+    return U, S, V
+  U, S, V = jaten._aten__linalg_svd(a, full_matrices=not some)
+  return U, S, jnp.matrix_transpose(V)
+
+@register_function(torch.cdist)
+def _cdist(x1, x2, p=2.0, compute_mode='use_mm_for_euclid_dist_if_necessary'):
+    return jaten._aten_cdist(x1, x2, p, compute_mode)
+
+@register_function(torch.lu)
+def lu(A, **kwargs):
+  lu,pivots,_ = jax.lax.linalg.lu(A)
+  # JAX pivots are offset by 1 compared to torch
+  _pivots = pivots + 1
+  info_shape = pivots.shape[:-1]
+  info = jnp.zeros(info_shape, dtype=mappings.t2j_dtype(torch.int32))
+  if kwargs['get_infos'] == True:
+    return lu, _pivots, info
+  return lu, _pivots
+
+@register_function(torch.lu_solve)
+def lu_solve(b, LU_data, LU_pivots, **kwargs):
+  # JAX pivots are offset by 1 compared to torch
+  _pivots = LU_pivots - 1
+  x = jax.scipy.linalg.lu_solve((LU_data, _pivots), b)
+  return x
+
+@register_function(torch.linalg.tensorsolve)
+def linalg_tensorsolve(A, b, dims=None):
+  # examples:
+  # A = torch.randn(2, 3, 6), b = torch.randn(3, 2)
+  # A = torch.randn(2, 3, 6), b = torch.randn(2, 3) -> torch.Size([3, 6])
+  # A = torch.randn(9, 2, 6, 3) b = torch.randn(6, 3) -> torch.Size([6, 3])
+  # A = torch.randn(9, 2, 3, 6) b = torch.randn(6, 3) -> torch.Size([3, 6])
+  # A = torch.randn(18, 6, 3) b = torch.randn(18) -> torch.Size([6, 3])
+  # A = torch.randn(3, 8, 4, 6) b = torch.randn(4, 6) -> torch.Size([4,6])
+  # A = torch.randn(3, 8, 1, 2, 2, 6) b = torch.randn(3, 4, 2) -> torch.Size([2, 2, 6])
+
+  # torch allows b to be shaped differently.
+  # especially when axes are moved using dims.
+  # ValueError: After moving axes to end, leading shape of a must match shape of b. got a.shape=(3, 2, 6), b.shape=(2, 3)
+  # So we are handling the moveaxis and forcing b's shape to match what jax expects
+  if dims is not None:
+    A = jnp.moveaxis(A, dims, len(dims) * (A.ndim - 1,))
+    dims = None
+  if A.shape[:b.ndim] != b.shape:
+    b = jnp.reshape(b, A.shape[:b.ndim])
+  return jnp.linalg.tensorsolve(A, b, axes=dims)
+
+
+@register_function(torch.nn.functional.linear)
+def functional_linear(self, weights, bias=None):
+  res = jnp.einsum("...a,ba->...b", self, weights)
+  if bias is not None:
+    res += bias
+  return res

torchax/ops/jtorchvision_nms.py

@@ -0,0 +1,245 @@
+"""
+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 torch
+from jax import lax
+import jax.numpy as jnp
+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
+  import torchvision
+  ops_registry.register_torch_dispatch_op(torch.ops.torchvision.nms, nms)
+except Exception:
+  pass