def testSortManually(self):
    ans = jnp.sort(onp.array([16, 15, 23, 42, 8, 4]))
    expected = onp.array([4, 8, 15, 16, 23, 42])
    self.assertAllClose(expected, ans, check_dtypes=True)
    a = onp.array([[1, 4], [3, 1]])
    ans = jnp.sort(a, axis=None)
    expected = onp.array([1, 1, 3, 4])
    self.assertAllClose(expected, ans, check_dtypes=True)
    a = onp.array([[1, 4], [3, 1]])
    ans = jnp.sort(a)
    expected = onp.array([[1, 4], [1, 3]])
    self.assertAllClose(expected, ans, check_dtypes=True)
    a = onp.array([[1, 4], [3, 1]])
    ans = jnp.sort(a, axis=0)
    expected = onp.array([[1, 1], [3, 4]])
    self.assertAllClose(expected, ans, check_dtypes=True)
def testSortManually(self):
    ans = jnp.sort(onp.array([16, 15, 23, 42, 8, 4]))
    expected = onp.array([4, 8, 15, 16, 23, 42])
    self.assertAllClose(expected, ans, check_dtypes=True)
    a = onp.array([[1, 4], [3, 1]])
    ans = jnp.sort(a, axis=None)
    expected = onp.array([1, 1, 3, 4])
    self.assertAllClose(expected, ans, check_dtypes=True)
    a = onp.array([[1, 4], [3, 1]])
    ans = jnp.sort(a)
    expected = onp.array([[1, 4], [1, 3]])
    self.assertAllClose(expected, ans, check_dtypes=True)
    a = onp.array([[1, 4], [3, 1]])
    ans = jnp.sort(a, axis=0)
    expected = onp.array([[1, 1], [3, 4]])
    self.assertAllClose(expected, ans, check_dtypes=True)