[sparse] ensure shapes are represented as tuples

hominluo · Nov 16, 2021 · 5c31e6d · 5c31e6d
1 parent 476ca94
commit 5c31e6d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/jax/experimental/sparse/ops.py b/jax/experimental/sparse/ops.py
@@ -730,7 +730,7 @@ def _safe_asarray(args):
     return map(_asarray_or_float0, args)
 
   def __init__(self, args, *, shape):
-    self.shape = shape
+    self.shape = tuple(shape)
 
   def __repr__(self):
     name = self.__class__.__name__

diff --git a/tests/sparse_test.py b/tests/sparse_test.py
@@ -1340,10 +1340,10 @@ def f(X, y):
 class SparseObjectTest(jtu.JaxTestCase):
   def test_repr(self):
     M = sparse.BCOO.fromdense(jnp.arange(5, dtype='float32'))
-    assert repr(M) == "BCOO(float32[5], nse=4)"
+    self.assertEqual(repr(M), "BCOO(float32[5], nse=4)")
 
-    M_invalid = sparse.BCOO(([], []), shape=100)
-    assert repr(M_invalid) == "BCOO(<invalid>)"
+    M_invalid = sparse.BCOO(([], []), shape=(100,))
+    self.assertEqual(repr(M_invalid), "BCOO(<invalid>)")
 
   @parameterized.named_parameters(
     {"testcase_name": "_{}".format(Obj.__name__), "Obj": Obj}