def testDisableNumpyRankPromotionBroadcasting(self):
    try:
        prev_flag = FLAGS.jax_numpy_rank_promotion
        FLAGS.jax_numpy_rank_promotion = 'allow'
        (jnp.ones(2) + jnp.ones((1, 2)))
    finally:
        FLAGS.jax_numpy_rank_promotion = prev_flag
    try:
        prev_flag = FLAGS.jax_numpy_rank_promotion
        FLAGS.jax_numpy_rank_promotion = 'raise'
        self.assertRaises(ValueError, (lambda : (jnp.ones(2) + jnp.ones((1, 2)))))
    finally:
        FLAGS.jax_numpy_rank_promotion = prev_flag
    try:
        prev_flag = FLAGS.jax_numpy_rank_promotion
        FLAGS.jax_numpy_rank_promotion = 'warn'
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always')
            (jnp.ones(2) + jnp.ones((1, 2)))
            assert (len(w) > 0)
            msg = str(w[(- 1)].message)
            expected_msg = 'Following NumPy automatic rank promotion for add on shapes (2,) (1, 2).'
            self.assertEqual(msg[:len(expected_msg)], expected_msg)
            prev_len = len(w)
            (jnp.ones(2) + 3)
            self.assertEqual(len(w), prev_len)
    finally:
        FLAGS.jax_numpy_rank_promotion = prev_flag
def testDisableNumpyRankPromotionBroadcasting(self):
    try:
        prev_flag = FLAGS.jax_numpy_rank_promotion
        FLAGS.jax_numpy_rank_promotion = 'allow'
        (jnp.ones(2) + jnp.ones((1, 2)))
    finally:
        FLAGS.jax_numpy_rank_promotion = prev_flag
    try:
        prev_flag = FLAGS.jax_numpy_rank_promotion
        FLAGS.jax_numpy_rank_promotion = 'raise'
        self.assertRaises(ValueError, (lambda : (jnp.ones(2) + jnp.ones((1, 2)))))
    finally:
        FLAGS.jax_numpy_rank_promotion = prev_flag
    try:
        prev_flag = FLAGS.jax_numpy_rank_promotion
        FLAGS.jax_numpy_rank_promotion = 'warn'
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always')
            (jnp.ones(2) + jnp.ones((1, 2)))
            assert (len(w) > 0)
            msg = str(w[(- 1)].message)
            expected_msg = 'Following NumPy automatic rank promotion for add on shapes (2,) (1, 2).'
            self.assertEqual(msg[:len(expected_msg)], expected_msg)
            prev_len = len(w)
            (jnp.ones(2) + 3)
            self.assertEqual(len(w), prev_len)
    finally:
        FLAGS.jax_numpy_rank_promotion = prev_flag