Python assert_axis_dimension Examples

Example #1

File: asserts_test.py Project: graingert/chex

 def test_assert_axis_dimension_fail(self):
   tensor = jnp.ones((3, 2, 7, 2))
   for i in range(-tensor.ndim, tensor.ndim):
     s = tensor.shape[i]
     with self.assertRaisesRegex(
         AssertionError, 'Expected tensor to have dimension'):
       asserts.assert_axis_dimension(tensor, axis=i, expected=s+1)

Example #2

File: asserts_test.py Project: graingert/chex

 def test_assert_axis_dimension_axis_invalid(self):
   tensor = jnp.ones((3, 2))
   for i in (2, -3):
     with self.assertRaisesRegex(AssertionError, 'not available'):
       asserts.assert_axis_dimension(tensor, axis=i, expected=1)

Example #3

File: asserts_test.py Project: stjordanis/chex

 def test_assert_axis_dimension_fail(self):
     tensor = jnp.ones((3, 2, 7, 2))
     for i, s in enumerate(tensor.shape):
         with self.assertRaisesRegex(AssertionError,
                                     'Expected tensor to have dimension'):
             asserts.assert_axis_dimension(tensor, axis=i, expected=s + 1)

Example #4

File: asserts_test.py Project: graingert/chex

 def test_assert_axis_dimension_pass(self):
   tensor = jnp.ones((3, 2, 7, 2))
   for i in range(-tensor.ndim, tensor.ndim):
     s = tensor.shape[i]
     asserts.assert_axis_dimension(tensor, axis=i, expected=s)

Example #5

File: asserts_test.py Project: stjordanis/chex

 def test_assert_axis_dimension_pass(self):
     tensor = jnp.ones((3, 2, 7, 2))
     for i, s in enumerate(tensor.shape):
         asserts.assert_axis_dimension(tensor, axis=i, expected=s)