Add torch.scatter

tests/test_core.py

@@ -1,9 +1,10 @@
 from functools import partial
 
 import jax.numpy as jnp
+import numpy as np
 import pytest
 import torch
-from jax import grad, jit, vmap
+from jax import grad, jit, random, vmap
 
 from torch2jax import t2j
 
@@ -249,6 +250,59 @@ def test_oneliners():
   t2j_function_test(lambda x: (x > 0.0).float(), [(3,)], tests=fb)
 
 
+def test_scatter():
+  # scatter
+  index = np.array([[0, 1, 2, 0, 2], [1, 0, 0, 2, 1]], dtype=np.int64)
+  samplers = [random.normal, lambda key, shape: index, random.normal]
+  tests = [forward_test, partial(backward_test, argnums=(0, 2))]
+  t2j_function_test(
+    lambda input, index, src: torch.scatter(input, 0, index, src),
+    [(3, 5), (2, 5), (2, 5)],
+    samplers=samplers,
+    atol=1e-6,
+    tests=tests,
+  )
+  # Disable gradient testing when index.shape != src.shape
+  # This is an existing problem of pytorch https://github.com/pytorch/pytorch/issues/27614
+  index = np.array([[0, 1, 2, 0, 2]], dtype=np.int64)
+  samplers = [random.normal, lambda key, shape: index, random.normal]
+  tests = [forward_test]
+  t2j_function_test(
+    lambda input, index, src: torch.scatter(input, 0, index, src),
+    [(3, 5), (1, 5), (2, 5)],
+    samplers=samplers,
+    atol=1e-6,
+    tests=tests,
+  )
+  index = np.array([[0, 1, 2, 0]], dtype=np.int64)
+  samplers = [random.normal, lambda key, shape: index, random.normal]
+  t2j_function_test(
+    lambda input, index, src: torch.scatter(input, 0, index, src),
+    [(3, 5), (1, 4), (2, 5)],
+    samplers=samplers,
+    atol=1e-6,
+    tests=tests,
+  )
+  index = np.array([[4, 2, 3], [3, 0, 4]], dtype=np.int64)
+  samplers = [random.normal, lambda key, shape: index, random.normal]
+  t2j_function_test(
+    lambda input, index, src: torch.scatter(input, 1, index, src),
+    [(3, 5), (2, 3), (3, 5)],
+    samplers=samplers,
+    atol=1e-6,
+    tests=tests,
+  )
+  index = np.array([[4, 2, 3], [3, 0, 4], [0, 1, 2]], dtype=np.int64)
+  samplers = [random.normal, lambda key, shape: index, random.normal]
+  t2j_function_test(
+    lambda input, index, src: torch.scatter(input, 1, index, src),
+    [(3, 5), (3, 3), (3, 5)],
+    samplers=samplers,
+    atol=1e-6,
+    tests=tests,
+  )
+
+
 def test_Tensor():
   with pytest.raises(ValueError):
     t2j(lambda: torch.Tensor([1, 2, 3]))()

tests/utils.py

@@ -70,9 +70,23 @@ def args_generator(shapes, samplers=None, rng=random.PRNGKey(123)):
     yield args
 
 
+def _arg2t(x):
+  if isinstance(x, jnp.ndarray):
+    # copy to avoid in-place modification
+    return j2t(x.copy())
+  elif isinstance(x, jnp.dtype):
+    return j2t(x)
+  elif isinstance(x, np.ndarray):
+    # we allow passing numpy arrays directly in test cases
+    # it is needed because torch requires int64 for some functions, while jax's default is int32, using numpy allows them to each make their own casts.
+    return torch.from_numpy(x)
+  else:
+    return x
+
+
 def forward_test(f, args, kwargs={}, test_jit=True, **assert_kwargs):
-  torch_args = jax.tree.map(lambda x: j2t(x.copy()) if isinstance(x, (jnp.ndarray, jnp.dtype)) else x, args)
-  torch_kwargs = jax.tree.map(lambda x: j2t(x) if isinstance(x, jnp.dtype) else x, kwargs)
+  torch_args = jax.tree.map(_arg2t, args)
+  torch_kwargs = jax.tree.map(_arg2t, kwargs)
   f_ = lambda *args, **kwargs: f(*args, **torch_kwargs)
   torch_output = f_(*torch_args)
   aac(t2j(f_)(*args), torch_output, **assert_kwargs)
@@ -85,8 +99,8 @@ def backward_test(f, args, kwargs={}, argnums=None, **assert_kwargs):
   argnums = argnums if argnums is not None else tuple(range(n_inputs))
   if n_inputs == 0 or len(argnums) == 0:
     return
-  torch_args = jax.tree.map(lambda x: j2t(x.copy()) if isinstance(x, (jnp.ndarray, jnp.dtype)) else x, args)
-  torch_kwargs = jax.tree.map(lambda x: j2t(x) if isinstance(x, jnp.dtype) else x, kwargs)
+  torch_args = jax.tree.map(_arg2t, args)
+  torch_kwargs = jax.tree.map(_arg2t, kwargs)
   # always reduce output to the mean of all elements
   f_ = lambda *args: torch.cat(list(map(lambda x: x.flatten(), tree_leaves(f(*args, **torch_kwargs))))).mean()
   for t2j_grad, torch_grad in zip(

torch2jax/__init__.py

@@ -1,6 +1,7 @@
 import copy
 import functools
 from contextlib import contextmanager
+from functools import partial
 from typing import Literal, Optional, Sequence, Tuple, Union
 
 import jax
@@ -86,10 +87,14 @@ def value(self):
   @value.setter
   def value(self, val):
     # See https://github.com/google/jax/issues/2115 re `isinstance(value, jnp.ndarray)`.
-    assert isinstance(val, jnp.ndarray) or isinstance(val, int) or isinstance(val, float), (
-      f"Tried to create Torchish with unsupported type: {type(val)}"
-    )
-    self._value = val if torch.is_grad_enabled() else jax.lax.stop_gradient(val)
+    if isinstance(val, jnp.ndarray):
+      self._value = val if torch.is_grad_enabled() else jax.lax.stop_gradient(val)
+    elif isinstance(val, (np.ndarray, int, float)):
+      # as jax is transparent to np.ndarray, we want the t2j functions to work on np.ndarray as well.
+      # stop_gradient doesn't affect np array or constants
+      self._value = val
+    else:
+      raise TypeError(f"Tried to create Torchish with unsupported type: {type(val)}")
 
   # In order for PyTorch to accept an object as one of its own and allow dynamic dispatch it must either subclass
   # `torch.Tensor` or have a `__torch_function__` method. We opt to take the method route. Dispatch logic is handled in
@@ -521,6 +526,43 @@ def _set_grad_enabled(mode):
   torch._C._set_grad_enabled(mode)
 
 
+def scatter_impl(input, dim, index, src, *, reduce=None):
+  # Code adopted from https://github.com/jax-ml/jax/issues/8487#issuecomment-1555311635
+  # in jax the dims other than dim for index & src should be equal,
+  # but torch allows index to be smaller, we pad it with a out-of-bound index to match the src.
+  out_of_range_idx = input.shape[dim]
+  padding = tuple((0, d2 - d1) for d1, d2 in zip(index.shape, src.shape))
+  index = jnp.pad(index, padding, constant_values=out_of_range_idx)
+  dnums = jax.lax.ScatterDimensionNumbers(
+    update_window_dims=(), inserted_window_dims=(0,), scatter_dims_to_operand_dims=(0,)
+  )
+  if reduce is None:
+    _scatter = jax.lax.scatter
+  elif reduce == "add":
+    _scatter = jax.lax.scatter_add
+  elif reduce == "multiply":
+    _scatter = jax.lax.scatter_mul
+
+  _scatter = partial(_scatter, dimension_numbers=dnums, mode="drop")
+  vmap_inner = partial(jax.vmap, in_axes=(0, 0, 0), out_axes=0)
+
+  for _ in range(len(input.shape) - 1):
+    _scatter = vmap_inner(_scatter)
+  swap = lambda x: jnp.swapaxes(x, dim, -1)
+  input, index, src = list(map(swap, (input, index, src)))
+  return swap(_scatter(input, jnp.expand_dims(index, axis=-1), src))
+
+
+@implements(torch.scatter, Torchish_member=True)
+def scatter(input, dim, index, src):
+  return scatter_impl(_v(input), dim, _v(index), _v(src), reduce=None)
+
+
+@implements(torch.scatter_add, Torchish_member=True)
+def scatter_add(input, dim, index, src):
+  return scatter_impl(_v(input), dim, _v(index), _v(src), reduce="add")
+
+
 @implements(torch.sort, out_kwarg=True, Torchish_member=True)
 def sort(input, dim=-1, descending=False, stable=False):
   return jnp.sort(_v(input), axis=dim, stable=stable, descending=descending)