test torch.transpose

samuela · Nov 7, 2023 · 13aa5d6 · 13aa5d6
1 parent 826eb15
commit 13aa5d6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/tests/test_all_the_things.py b/tests/test_all_the_things.py
@@ -122,6 +122,10 @@ def test_oneliners():
 
   t2j_function_test(lambda x: x.expand(5, -1, -1), [(1, 3, 2)])
 
+  t2j_function_test(lambda x: torch.transpose(x, 0, 1), [(2, 3)])
+  t2j_function_test(lambda x: torch.transpose(x, 0, 2), [(2, 3, 5)])
+  t2j_function_test(lambda x: torch.transpose(x, 2, 1), [(2, 3, 5)])
+
 
 def test_detach():
   t2j_function_test(lambda x: x.detach() ** 2, [()])

diff --git a/torch2jax/__init__.py b/torch2jax/__init__.py
@@ -168,7 +168,6 @@ def fn(*args, **kwargs):
 connect(torch.sqrt, jnp.sqrt)
 connect(torch.sum, jnp.sum)
 connect(torch.tanh, jnp.tanh)
-# TODO: test this... it should need dont_coerce_argnums
 connect(torch.transpose, jnp.swapaxes)
 
 connect(torch.Tensor.mul, jnp.multiply)