text: import contextlib
  import sys
  import numpy as np
  import random
  import re
  import threading
  import gc
  
  from numba.core.errors import TypingError
  from numba import njit
  from numba.core import types, utils, config
  from numba.tests.support import MemoryLeakMixin, TestCase, tag
  import unittest
  
  
  nrtjit = njit(_nrt=True, nogil=True)
  
  
  def np_concatenate1(a, b, c):
      return np.concatenate((a, b, c))
  
  def np_concatenate2(a, b, c, axis):
      return np.concatenate((a, b, c), axis=axis)
  
  def np_stack1(a, b, c):
      return np.stack((a, b, c))
  
  def np_stack2(a, b, c, axis):
      return np.stack((a, b, c), axis=axis)
  
  def np_hstack(a, b, c):
      return np.hstack((a, b, c))
  
  def np_vstack(a, b, c):
      return np.vstack((a, b, c))
  
  def np_dstack(a, b, c):
      return np.dstack((a, b, c))
  
  def np_column_stack(a, b, c):
      return np.column_stack((a, b, c))
  
  
  class BaseTest(TestCase):
  
      def check_outputs(self, pyfunc, argslist, exact=True):
          cfunc = nrtjit(pyfunc)
          for args in argslist:
              expected = pyfunc(*args)
              ret = cfunc(*args)
              self.assertEqual(ret.size, expected.size)
              self.assertEqual(ret.dtype, expected.dtype)
              self.assertStridesEqual(ret, expected)
              if exact:
                  np.testing.assert_equal(expected, ret)
              else:
                  np.testing.assert_allclose(expected, ret)
  
  
  class NrtRefCtTest(MemoryLeakMixin):
      def assert_array_nrt_refct(self, arr, expect):
          self.assertEqual(arr.base.refcount, expect)
  
  
  class TestDynArray(NrtRefCtTest, TestCase):
  
      def test_empty_0d(self):
          @nrtjit
          def foo():
              arr = np.empty(())
              arr[()] = 42
              return arr
  
          arr = foo()
          self.assert_array_nrt_refct(arr, 1)
          np.testing.assert_equal(42, arr)
          self.assertEqual(arr.size, 1)
          self.assertEqual(arr.shape, ())
          self.assertEqual(arr.dtype, np.dtype(np.float64))
          self.assertEqual(arr.strides, ())
          arr.fill(123)  # test writability
          np.testing.assert_equal(123, arr)
          del arr
  
      def test_empty_1d(self):
          @nrtjit
          def foo(n):
              arr = np.empty(n)
              for i in range(n):
                  arr[i] = i
  
              return arr
  
          n = 3
          arr = foo(n)
          self.assert_array_nrt_refct(arr, 1)
          np.testing.assert_equal(np.arange(n), arr)
          self.assertEqual(arr.size, n)
          self.assertEqual(arr.shape, (n,))
          self.assertEqual(arr.dtype, np.dtype(np.float64))
          self.assertEqual(arr.strides, (np.dtype(np.float64).itemsize,))
          arr.fill(123)  # test writability
          np.testing.assert_equal(123, arr)
          del arr
  
      def test_empty_2d(self):
          def pyfunc(m, n):
              arr = np.empty((m, n), np.int32)
              for i in range(m):
                  for j in range(n):
                      arr[i, j] = i + j
  
              return arr
  
          cfunc = nrtjit(pyfunc)
          m = 4
          n = 3
          expected_arr = pyfunc(m, n)
          got_arr = cfunc(m, n)
          self.assert_array_nrt_refct(got_arr, 1)
          np.testing.assert_equal(expected_arr, got_arr)
  
          self.assertEqual(expected_arr.size, got_arr.size)
          self.assertEqual(expected_arr.shape, got_arr.shape)
          self.assertEqual(expected_arr.strides, got_arr.strides)
  
          del got_arr
  
      def test_empty_3d(self):
          def pyfunc(m, n, p):
              arr = np.empty((m, n, p), np.int32)
              for i in range(m):
                  for j in range(n):
                      for k in range(p):
                          arr[i, j, k] = i + j + k
  
              return arr
  
          cfunc = nrtjit(pyfunc)
          m = 4
          n = 3
          p = 2
          expected_arr = pyfunc(m, n, p)
          got_arr = cfunc(m, n, p)
          self.assert_array_nrt_refct(got_arr, 1)
          np.testing.assert_equal(expected_arr, got_arr)
  
          self.assertEqual(expected_arr.size, got_arr.size)
          self.assertEqual(expected_arr.shape, got_arr.shape)
          self.assertEqual(expected_arr.strides, got_arr.strides)
  
          del got_arr
  
      def test_empty_2d_sliced(self):
          def pyfunc(m, n, p):
              arr = np.empty((m, n), np.int32)
              for i in range(m):
                  for j in range(n):
                      arr[i, j] = i + j
  
              return arr[p]
  
          cfunc = nrtjit(pyfunc)
          m = 4
          n = 3
          p = 2
          expected_arr = pyfunc(m, n, p)
          got_arr = cfunc(m, n, p)
          self.assert_array_nrt_refct(got_arr, 1)
          np.testing.assert_equal(expected_arr, got_arr)
  
          self.assertEqual(expected_arr.size, got_arr.size)
          self.assertEqual(expected_arr.shape, got_arr.shape)
          self.assertEqual(expected_arr.strides, got_arr.strides)
  
          del got_arr
  
      def test_return_global_array(self):
          y = np.ones(4, dtype=np.float32)
          initrefct = sys.getrefcount(y)
  
          def return_external_array():
              return y
  
          cfunc = nrtjit(return_external_array)
          out = cfunc()
  
          # out reference by cfunc
          self.assertEqual(initrefct + 1, sys.getrefcount(y))
  
          np.testing.assert_equal(y, out)
          np.testing.assert_equal(y, np.ones(4, dtype=np.float32))
          np.testing.assert_equal(out, np.ones(4, dtype=np.float32))
  
          del out
          gc.collect()
          # out is only referenced by cfunc
          self.assertEqual(initrefct + 1, sys.getrefcount(y))
  
          del cfunc
          gc.collect()
          # y is no longer referenced by cfunc
          self.assertEqual(initrefct, sys.getrefcount(y))
  
      def test_return_global_array_sliced(self):
          y = np.ones(4, dtype=np.float32)
  
          def return_external_array():
              return y[2:]
  
          cfunc = nrtjit(return_external_array)
          out = cfunc()
          self.assertIsNone(out.base)
  
          yy = y[2:]
          np.testing.assert_equal(yy, out)
          np.testing.assert_equal(yy, np.ones(2, dtype=np.float32))
          np.testing.assert_equal(out, np.ones(2, dtype=np.float32))
  
      def test_array_pass_through(self):
          def pyfunc(y):
              return y
  
          arr = np.ones(4, dtype=np.float32)
  
          cfunc = nrtjit(pyfunc)
          expected = cfunc(arr)
          got = pyfunc(arr)
  
          np.testing.assert_equal(expected, arr)
          np.testing.assert_equal(expected, got)
          self.assertIs(expected, arr)
          self.assertIs(expected, got)
  
      def test_array_pass_through_sliced(self):
          def pyfunc(y):
              return y[y.size // 2:]
  
          arr = np.ones(4, dtype=np.float32)
  
          initrefct = sys.getrefcount(arr)
  
          cfunc = nrtjit(pyfunc)
          got = cfunc(arr)
          self.assertEqual(initrefct + 1, sys.getrefcount(arr))
          expected = pyfunc(arr)
          self.assertEqual(initrefct + 2, sys.getrefcount(arr))
  
          np.testing.assert_equal(expected, arr[arr.size // 2])
          np.testing.assert_equal(expected, got)
  
          del expected
          self.assertEqual(initrefct + 1, sys.getrefcount(arr))
          del got
          self.assertEqual(initrefct, sys.getrefcount(arr))
  
      def test_ufunc_with_allocated_output(self):
  
          def pyfunc(a, b):
              out = np.empty(a.shape)
              np.add(a, b, out)
              return out
  
          cfunc = nrtjit(pyfunc)
  
          # 1D case
          arr_a = np.random.random(10)
          arr_b = np.random.random(10)
  
          np.testing.assert_equal(pyfunc(arr_a, arr_b),
                                  cfunc(arr_a, arr_b))
  
          self.assert_array_nrt_refct(cfunc(arr_a, arr_b), 1)
  
          # 2D case
          arr_a = np.random.random(10).reshape(2, 5)
          arr_b = np.random.random(10).reshape(2, 5)
  
          np.testing.assert_equal(pyfunc(arr_a, arr_b),
                                  cfunc(arr_a, arr_b))
  
          self.assert_array_nrt_refct(cfunc(arr_a, arr_b), 1)
  
          # 3D case
          arr_a = np.random.random(70).reshape(2, 5, 7)
          arr_b = np.random.random(70).reshape(2, 5, 7)
  
          np.testing.assert_equal(pyfunc(arr_a, arr_b),
                                  cfunc(arr_a, arr_b))
  
          self.assert_array_nrt_refct(cfunc(arr_a, arr_b), 1)
  
      def test_allocation_mt(self):
          """
          This test exercises the array allocation in multithreaded usecase.
          This stress the freelist inside NRT.
          """
  
          def pyfunc(inp):
              out = np.empty(inp.size)
  
              # Zero fill
              for i in range(out.size):
                  out[i] = 0
  
              for i in range(inp[0]):
                  # Allocate inside a loop
                  tmp = np.empty(inp.size)
                  # Write to tmp
                  for j in range(tmp.size):
                      tmp[j] = inp[j]
                  # out = tmp + i
                  for j in range(tmp.size):
                      out[j] += tmp[j] + i
  
              return out
  
          cfunc = nrtjit(pyfunc)
          size = 10  # small array size so that the computation is short
          arr = np.random.randint(1, 10, size)
          frozen_arr = arr.copy()
  
          np.testing.assert_equal(pyfunc(arr), cfunc(arr))
          # Ensure we did not modify the input
          np.testing.assert_equal(frozen_arr, arr)
  
          workers = []
          inputs = []
          outputs = []
  
          # Make wrapper to store the output
          def wrapped(inp, out):
              out[:] = cfunc(inp)
  
          # Create a lot of worker threads to create contention
          for i in range(100):
              arr = np.random.randint(1, 10, size)
              out = np.empty_like(arr)
              thread = threading.Thread(target=wrapped,
                                        args=(arr, out),
                                        name="worker{0}".format(i))
              workers.append(thread)
              inputs.append(arr)
              outputs.append(out)
  
          # Launch worker threads
          for thread in workers:
              thread.start()
  
          # Join worker threads
          for thread in workers:
              thread.join()
  
          # Check result
          for inp, out in zip(inputs, outputs):
              np.testing.assert_equal(pyfunc(inp), out)
  
      def test_refct_mt(self):
          """
          This test exercises the refct in multithreaded code
          """
  
          def pyfunc(n, inp):
              out = np.empty(inp.size)
              for i in range(out.size):
                  out[i] = inp[i] + 1
              # Use swap to trigger many refct ops
              for i in range(n):
                  out, inp = inp, out
              return out
  
          cfunc = nrtjit(pyfunc)
          size = 10
          input = np.arange(size, dtype=float)
          expected_refct = sys.getrefcount(input)
          swapct = random.randrange(1000)
          expected = pyfunc(swapct, input)
          np.testing.assert_equal(expected, cfunc(swapct, input))
          # The following checks can discover a reference count error
          del expected
          self.assertEqual(expected_refct, sys.getrefcount(input))
  
          workers = []
          outputs = []
          swapcts = []
  
          # Make wrapper to store the output
          def wrapped(n, input, out):
              out[:] = cfunc(n, input)
  
          # Create worker threads
          for i in range(100):
              out = np.empty(size)
              # All thread shares the same input
              swapct = random.randrange(1000)
              thread = threading.Thread(target=wrapped,
                                        args=(swapct, input, out),
                                        name="worker{0}".format(i))
              workers.append(thread)
              outputs.append(out)
              swapcts.append(swapct)
  
          # Launch worker threads
          for thread in workers:
              thread.start()
  
          # Join worker threads
          for thread in workers:
              thread.join()
  
          # Check result
          for swapct, out in zip(swapcts, outputs):
              np.testing.assert_equal(pyfunc(swapct, input), out)
  
          del outputs, workers
          # The following checks can discover a reference count error
          self.assertEqual(expected_refct, sys.getrefcount(input))
  
      def test_swap(self):
  
          def pyfunc(x, y, t):
              """Swap array x and y for t number of times
              """
              for i in range(t):
                  x, y = y, x
  
              return x, y
  
  
          cfunc = nrtjit(pyfunc)
  
          x = np.random.random(100)
          y = np.random.random(100)
  
          t = 100
  
          initrefct = sys.getrefcount(x), sys.getrefcount(y)
          expect, got = pyfunc(x, y, t), cfunc(x, y, t)
          self.assertIsNone(got[0].base)
          self.assertIsNone(got[1].base)
          np.testing.assert_equal(expect, got)
          del expect, got
          self.assertEqual(initrefct, (sys.getrefcount(x), sys.getrefcount(y)))
  
      def test_return_tuple_of_array(self):
  
          def pyfunc(x):
              y = np.empty(x.size)
              for i in range(y.size):
                  y[i] = x[i] + 1
              return x, y
  
          cfunc = nrtjit(pyfunc)
  
          x = np.random.random(5)
          initrefct = sys.getrefcount(x)
          expected_x, expected_y = pyfunc(x)
          got_x, got_y = cfunc(x)
          self.assertIs(x, expected_x)
          self.assertIs(x, got_x)
          np.testing.assert_equal(expected_x, got_x)
          np.testing.assert_equal(expected_y, got_y)
          del expected_x, got_x
          self.assertEqual(initrefct, sys.getrefcount(x))
  
          self.assertEqual(sys.getrefcount(expected_y), sys.getrefcount(got_y))
  
      def test_return_tuple_of_array_created(self):
  
          def pyfunc(x):
              y = np.empty(x.size)
              for i in range(y.size):
                  y[i] = x[i] + 1
              out = y, y
              return out
  
          cfunc = nrtjit(pyfunc)
  
          x = np.random.random(5)
          expected_x, expected_y = pyfunc(x)
          got_x, got_y = cfunc(x)
          np.testing.assert_equal(expected_x, got_x)
          np.testing.assert_equal(expected_y, got_y)
          # getrefcount owns 1, got_y owns 1
          self.assertEqual(2, sys.getrefcount(got_y))
          # getrefcount owns 1, got_y owns 1
          self.assertEqual(2, sys.getrefcount(got_y))
  
      def test_issue_with_return_leak(self):
          """
          Dispatcher returns a new reference.
          It need to workaround it for now.
          """
          @nrtjit
          def inner(out):
              return out
  
          def pyfunc(x):
              return inner(x)
  
          cfunc = nrtjit(pyfunc)
  
          arr = np.arange(10)
          old_refct = sys.getrefcount(arr)
  
          self.assertEqual(old_refct, sys.getrefcount(pyfunc(arr)))
          self.assertEqual(old_refct, sys.getrefcount(cfunc(arr)))
          self.assertEqual(old_refct, sys.getrefcount(arr))
  
  
  class ConstructorBaseTest(NrtRefCtTest):
  
      def check_0d(self, pyfunc):
          cfunc = nrtjit(pyfunc)
          expected = pyfunc()
          ret = cfunc()
          self.assert_array_nrt_refct(ret, 1)
          self.assertEqual(ret.size, expected.size)
          self.assertEqual(ret.shape, expected.shape)
          self.assertEqual(ret.dtype, expected.dtype)
          self.assertEqual(ret.strides, expected.strides)
          self.check_result_value(ret, expected)
          # test writability
          expected = np.empty_like(ret) # np.full_like was not added until Numpy 1.8
          expected.fill(123)
          ret.fill(123)
          np.testing.assert_equal(ret, expected)
  
      def check_1d(self, pyfunc):
          cfunc = nrtjit(pyfunc)
          n = 3
          expected = pyfunc(n)
          ret = cfunc(n)
          self.assert_array_nrt_refct(ret, 1)
          self.assertEqual(ret.size, expected.size)
          self.assertEqual(ret.shape, expected.shape)
          self.assertEqual(ret.dtype, expected.dtype)
          self.assertEqual(ret.strides, expected.strides)
          self.check_result_value(ret, expected)
          # test writability
          expected = np.empty_like(ret) # np.full_like was not added until Numpy 1.8
          expected.fill(123)
          ret.fill(123)
          np.testing.assert_equal(ret, expected)
          # errors
          with self.assertRaises(ValueError) as cm:
              cfunc(-1)
          self.assertEqual(str(cm.exception), "negative dimensions not allowed")
  
      def check_2d(self, pyfunc):
          cfunc = nrtjit(pyfunc)
          m, n = 2, 3
          expected = pyfunc(m, n)
          ret = cfunc(m, n)
          self.assert_array_nrt_refct(ret, 1)
          self.assertEqual(ret.size, expected.size)
          self.assertEqual(ret.shape, expected.shape)
          self.assertEqual(ret.dtype, expected.dtype)
          self.assertEqual(ret.strides, expected.strides)
          self.check_result_value(ret, expected)
          # test writability
          expected = np.empty_like(ret)  # np.full_like was not added until Numpy 1.8
          expected.fill(123)
          ret.fill(123)
          np.testing.assert_equal(ret, expected)
          # errors
          with self.assertRaises(ValueError) as cm:
              cfunc(2, -1)
          self.assertEqual(str(cm.exception), "negative dimensions not allowed")
  
      def check_alloc_size(self, pyfunc):
          """Checks that pyfunc will error, not segfaulting due to array size."""
          cfunc = nrtjit(pyfunc)
          with self.assertRaises(ValueError) as e:
              cfunc()
          self.assertIn(
              "array is too big",
              str(e.exception)
          )
  
  
  class TestNdZeros(ConstructorBaseTest, TestCase):
  
      def setUp(self):
          super(TestNdZeros, self).setUp()
          self.pyfunc = np.zeros
  
      def check_result_value(self, ret, expected):
          np.testing.assert_equal(ret, expected)
  
      def test_0d(self):
          pyfunc = self.pyfunc
          def func():
              return pyfunc(())
          self.check_0d(func)
  
      def test_1d(self):
          pyfunc = self.pyfunc
          def func(n):
              return pyfunc(n)
          self.check_1d(func)
  
      def test_1d_dtype(self):
          pyfunc = self.pyfunc
          def func(n):
              return pyfunc(n, np.int32)
          self.check_1d(func)
  
      def test_1d_dtype_instance(self):
          # dtype as numpy dtype, not as scalar class
          pyfunc = self.pyfunc
          _dtype = np.dtype('int32')
          def func(n):
              return pyfunc(n, _dtype)
          self.check_1d(func)
  
      def test_1d_dtype_str(self):
          pyfunc = self.pyfunc
          _dtype = 'int32'
          def func(n):
              return pyfunc(n, _dtype)
          self.check_1d(func)
  
          def func(n):
              return pyfunc(n, 'complex128')
          self.check_1d(func)
  
      def test_1d_dtype_str_alternative_spelling(self):
          # like test_1d_dtype_str but using the shorthand type spellings
          pyfunc = self.pyfunc
          _dtype = 'i4'
          def func(n):
              return pyfunc(n, _dtype)
          self.check_1d(func)
  
          def func(n):
              return pyfunc(n, 'c8')
          self.check_1d(func)
  
      def test_1d_dtype_str_structured_dtype(self):
          # test_1d_dtype_str but using a structured dtype
          pyfunc = self.pyfunc
          _dtype = "i4, (2,3)f8"
          def func(n):
              return pyfunc(n, _dtype)
          self.check_1d(func)
  
      def test_1d_dtype_non_const_str(self):
          pyfunc = self.pyfunc
  
          @njit
          def func(n, dt):
              return pyfunc(n, dt)
  
          with self.assertRaises(TypingError) as raises:
              func(5, 'int32')
  
          excstr = str(raises.exception)
          self.assertIn('No match', excstr)
          restr = r'\b{}\(int.*?, unicode_type\)\B'
          regex = re.compile(restr.format(pyfunc.__name__))
          self.assertRegex(excstr, regex)
  
      def test_1d_dtype_invalid_str(self):
          pyfunc = self.pyfunc
  
          @njit
          def func(n):
              return pyfunc(n, 'ABCDEF')
  
          with self.assertRaises(TypingError) as raises:
              func(5)
  
          excstr = str(raises.exception)
          self.assertIn("Invalid NumPy dtype specified: 'ABCDEF'", excstr)
  
      def test_2d(self):
          pyfunc = self.pyfunc
          def func(m, n):
              return pyfunc((m, n))
          self.check_2d(func)
  
      def test_2d_shape_dtypes(self):
          # Test for issue #4575
          pyfunc = self.pyfunc
          def func1(m, n):
              return pyfunc((np.int16(m), np.int32(n)))
          self.check_2d(func1)
          # Using a 64-bit value checks that 32 bit systems will downcast to intp
          def func2(m, n):
              return pyfunc((np.int64(m), np.int8(n)))
          self.check_2d(func2)
          # Make sure an error is thrown if we can't downcast safely
          if config.IS_32BITS:
              cfunc = nrtjit(lambda m, n: pyfunc((m, n)))
              with self.assertRaises(ValueError):
                  cfunc(np.int64(1 << (32 - 1)), 1)
  
      def test_2d_dtype_kwarg(self):
          pyfunc = self.pyfunc
          def func(m, n):
              return pyfunc((m, n), dtype=np.complex64)
          self.check_2d(func)
  
      def test_2d_dtype_str_kwarg(self):
          pyfunc = self.pyfunc
          def func(m, n):
              return pyfunc((m, n), dtype='complex64')
          self.check_2d(func)
  
      def test_2d_dtype_str_kwarg_alternative_spelling(self):
          # as test_2d_dtype_str_kwarg but with the numpy shorthand type spelling
          pyfunc = self.pyfunc
          def func(m, n):
              return pyfunc((m, n), dtype='c8')
          self.check_2d(func)
  
      def test_alloc_size(self):
          pyfunc = self.pyfunc
          width = types.intp.bitwidth
          def gen_func(shape, dtype):
              return lambda : pyfunc(shape, dtype)
          # Under these values numba will segfault, but thats another issue
          self.check_alloc_size(gen_func(1 << width - 2, np.intp))
          self.check_alloc_size(gen_func((1 << width - 8, 64), np.intp))
  
  
  class TestNdOnes(TestNdZeros):
  
      def setUp(self):
          super(TestNdOnes, self).setUp()
          self.pyfunc = np.ones
  
      @unittest.expectedFailure
      def test_1d_dtype_str_structured_dtype(self):
          super().test_1d_dtype_str_structured_dtype()
  
  
  class TestNdFull(ConstructorBaseTest, TestCase):
  
      def check_result_value(self, ret, expected):
          np.testing.assert_equal(ret, expected)
  
      def test_0d(self):
          def func():
              return np.full((), 4.5)
          self.check_0d(func)
  
      def test_1d(self):
          def func(n):
              return np.full(n, 4.5)
          self.check_1d(func)
  
      def test_1d_dtype(self):
          def func(n):
              return np.full(n, 4.5, np.bool_)
          self.check_1d(func)
  
      def test_1d_dtype_instance(self):
          dtype = np.dtype('bool')
          def func(n):
              return np.full(n, 4.5, dtype)
          self.check_1d(func)
  
      def test_1d_dtype_str(self):
          def func(n):
              return np.full(n, 4.5, 'bool_')
          self.check_1d(func)
  
      def test_1d_dtype_str_alternative_spelling(self):
          # like test_1d_dtype_str but using the shorthand type spelling
          def func(n):
              return np.full(n, 4.5, '?')
          self.check_1d(func)
  
      def test_1d_dtype_non_const_str(self):
  
          @njit
          def func(n, fv, dt):
              return np.full(n, fv, dt)
  
          with self.assertRaises(TypingError) as raises:
              func((5,), 4.5, 'int32')
  
          excstr = str(raises.exception)
          self.assertIn('No match', excstr)
          restr = r'\bfull\(UniTuple\(int.*? x 1\), float64, unicode_type\)\B'
          regex = re.compile(restr)
          self.assertRegex(excstr, regex)
  
      def test_1d_dtype_invalid_str(self):
  
          @njit
          def func(n, fv):
              return np.full(n, fv, 'ABCDEF')
  
          with self.assertRaises(TypingError) as raises:
              func(np.ones(4), 4.5)
  
          excstr = str(raises.exception)
          self.assertIn("Invalid NumPy dtype specified: 'ABCDEF'", excstr)
  
      def test_2d(self):
          def func(m, n):
              return np.full((m, n), 4.5)
          self.check_2d(func)
  
      def test_2d_dtype_kwarg(self):
          def func(m, n):
              return np.full((m, n), 1 + 4.5j, dtype=np.complex64)
          self.check_2d(func)
  
      def test_2d_dtype_from_type(self):
          # tests issue #2862
          def func(m, n):
              return np.full((m, n), np.int32(1))
          self.check_2d(func)
  
          # Complex uses `.real`, imaginary part dropped
          def func(m, n):
              return np.full((m, n), np.complex128(1))
          self.check_2d(func)
  
          # and that if a dtype is specified, this influences the return type
          def func(m, n):
              return np.full((m, n), 1, dtype=np.int8)
          self.check_2d(func)
  
      def test_2d_shape_dtypes(self):
          # Test for issue #4575
          def func1(m, n):
              return np.full((np.int16(m), np.int32(n)), 4.5)
          self.check_2d(func1)
          # Using a 64-bit value checks that 32 bit systems will downcast to intp
          def func2(m, n):
              return np.full((np.int64(m), np.int8(n)), 4.5)
          self.check_2d(func2)
          # Make sure an error is thrown if we can't downcast safely
          if config.IS_32BITS:
              cfunc = nrtjit(lambda m, n: np.full((m, n), 4.5))
              with self.assertRaises(ValueError):
                  cfunc(np.int64(1 << (32 - 1)), 1)
  
      def test_alloc_size(self):
          width = types.intp.bitwidth
          def gen_func(shape, value):
              return lambda : np.full(shape, value)
          # Under these values numba will segfault, but thats another issue
          self.check_alloc_size(gen_func(1 << width - 2, 1))
          self.check_alloc_size(gen_func((1 << width - 8, 64), 1))
  
  
  class ConstructorLikeBaseTest(object):
  
      def mutate_array(self, arr):
          try:
              arr.fill(42)
          except (TypeError, ValueError):
              # Try something else (e.g. Numpy 1.6 with structured dtypes)
              fill_value = b'x' * arr.dtype.itemsize
              arr.fill(fill_value)
  
      def check_like(self, pyfunc, dtype):
          def check_arr(arr):
              expected = pyfunc(arr)
              ret = cfunc(arr)
              self.assertEqual(ret.size, expected.size)
              self.assertEqual(ret.dtype, expected.dtype)
              self.assertStridesEqual(ret, expected)
              self.check_result_value(ret, expected)
              # test writability
              self.mutate_array(ret)
              self.mutate_array(expected)
              np.testing.assert_equal(ret, expected)
  
          orig = np.linspace(0, 5, 6).astype(dtype)
          cfunc = nrtjit(pyfunc)
  
          for shape in (6, (2, 3), (1, 2, 3), (3, 1, 2), ()):
              if shape == ():
                  arr = orig[-1:].reshape(())
              else:
                  arr = orig.reshape(shape)
              check_arr(arr)
              # Non-contiguous array
              if arr.ndim > 0:
                  check_arr(arr[::2])
              # Check new array doesn't inherit readonly flag
              arr.flags['WRITEABLE'] = False
              # verify read-only
              with self.assertRaises(ValueError):
                  arr[0] = 1
              check_arr(arr)
  
          # Scalar argument => should produce a 0-d array
          check_arr(orig[0])
  
  
  class TestNdEmptyLike(ConstructorLikeBaseTest, TestCase):
  
      def setUp(self):
          super(TestNdEmptyLike, self).setUp()
          self.pyfunc = np.empty_like
  
      def check_result_value(self, ret, expected):
          pass
  
      def test_like(self):
          pyfunc = self.pyfunc
          def func(arr):
              return pyfunc(arr)
          self.check_like(func, np.float64)
  
      def test_like_structured(self):
          dtype = np.dtype([('a', np.int16), ('b', np.float32)])
          pyfunc = self.pyfunc
          def func(arr):
              return pyfunc(arr)
          self.check_like(func, dtype)
  
      def test_like_dtype(self):
          pyfunc = self.pyfunc
          def func(arr):
              return pyfunc(arr, np.int32)
          self.check_like(func, np.float64)
  
      def test_like_dtype_instance(self):
          dtype = np.dtype('int32')
          pyfunc = self.pyfunc
          def func(arr):
              return pyfunc(arr, dtype)
          self.check_like(func, np.float64)
  
      def test_like_dtype_structured(self):
          dtype = np.dtype([('a', np.int16), ('b', np.float32)])
          pyfunc = self.pyfunc
          def func(arr):
              return pyfunc(arr, dtype)
          self.check_like(func, np.float64)
  
      def test_like_dtype_kwarg(self):
          pyfunc = self.pyfunc
          def func(arr):
              return pyfunc(arr, dtype=np.int32)
          self.check_like(func, np.float64)
  
      def test_like_dtype_str_kwarg(self):
          pyfunc = self.pyfunc
          def func(arr):
              return pyfunc(arr, dtype='int32')
          self.check_like(func, np.float64)
  
      def test_like_dtype_str_kwarg_alternative_spelling(self):
          pyfunc = self.pyfunc
          def func(arr):
              return pyfunc(arr, dtype='i4')
          self.check_like(func, np.float64)
  
      def test_like_dtype_non_const_str(self):
          pyfunc = self.pyfunc
  
          @njit
          def func(n, dt):
              return pyfunc(n, dt)
  
          with self.assertRaises(TypingError) as raises:
              func(np.ones(4), 'int32')
  
          excstr = str(raises.exception)
  
          self.assertIn('No match', excstr)
          self.assertIn(
              '{}(array(float64, 1d, C), unicode_type)'.format(pyfunc.__name__),
              excstr)
  
      def test_like_dtype_invalid_str(self):
          pyfunc = self.pyfunc
  
          @njit
          def func(n):
              return pyfunc(n, 'ABCDEF')
  
          with self.assertRaises(TypingError) as raises:
              func(np.ones(4))
  
          excstr = str(raises.exception)
          self.assertIn("Invalid NumPy dtype specified: 'ABCDEF'", excstr)
  
  
  class TestNdZerosLike(TestNdEmptyLike):
  
      def setUp(self):
          super(TestNdZerosLike, self).setUp()
          self.pyfunc = np.zeros_like
  
      def check_result_value(self, ret, expected):
          np.testing.assert_equal(ret, expected)
  
      def test_like_structured(self):
          super(TestNdZerosLike, self).test_like_structured()
  
      def test_like_dtype_structured(self):
          super(TestNdZerosLike, self).test_like_dtype_structured()
  
  
  class TestNdOnesLike(TestNdZerosLike):
  
      def setUp(self):
          super(TestNdOnesLike, self).setUp()
          self.pyfunc = np.ones_like
          self.expected_value = 1
  
      # Not supported yet.
  
      @unittest.expectedFailure
      def test_like_structured(self):
          super(TestNdOnesLike, self).test_like_structured()
  
      @unittest.expectedFailure
      def test_like_dtype_structured(self):
          super(TestNdOnesLike, self).test_like_dtype_structured()
  
  
  class TestNdFullLike(ConstructorLikeBaseTest, TestCase):
  
      def check_result_value(self, ret, expected):
          np.testing.assert_equal(ret, expected)
  
      def test_like(self):
          def func(arr):
              return np.full_like(arr, 3.5)
          self.check_like(func, np.float64)
  
      # Not supported yet.
      @unittest.expectedFailure
      def test_like_structured(self):
          dtype = np.dtype([('a', np.int16), ('b', np.float32)])
          def func(arr):
              return np.full_like(arr, 4.5)
          self.check_like(func, dtype)
  
      def test_like_dtype(self):
          def func(arr):
              return np.full_like(arr, 4.5, np.bool_)
          self.check_like(func, np.float64)
  
      def test_like_dtype_instance(self):
          dtype = np.dtype('bool')
          def func(arr):
              return np.full_like(arr, 4.5, dtype)
          self.check_like(func, np.float64)
  
      def test_like_dtype_kwarg(self):
          def func(arr):
              return np.full_like(arr, 4.5, dtype=np.bool_)
          self.check_like(func, np.float64)
  
      def test_like_dtype_str_kwarg(self):
          def func(arr):
              return np.full_like(arr, 4.5, 'bool_')
          self.check_like(func, np.float64)
  
      def test_like_dtype_str_kwarg_alternative_spelling(self):
          def func(arr):
              return np.full_like(arr, 4.5, dtype='?')
          self.check_like(func, np.float64)
  
      def test_like_dtype_non_const_str_kwarg(self):
  
          @njit
          def func(arr, fv, dt):
              return np.full_like(arr, fv, dt)
  
          with self.assertRaises(TypingError) as raises:
              func(np.ones(3,), 4.5, 'int32')
  
          excstr = str(raises.exception)
          self.assertIn('No match', excstr)
          self.assertIn('full_like(array(float64, 1d, C), float64, unicode_type)',
                        excstr)
  
      def test_like_dtype_invalid_str(self):
  
          @njit
          def func(arr, fv):
              return np.full_like(arr, fv, "ABCDEF")
  
          with self.assertRaises(TypingError) as raises:
              func(np.ones(4), 3.4)
  
          excstr = str(raises.exception)
          self.assertIn("Invalid NumPy dtype specified: 'ABCDEF'", excstr)
  
  
  class TestNdIdentity(BaseTest):
  
      def check_identity(self, pyfunc):
          self.check_outputs(pyfunc, [(3,)])
  
      def test_identity(self):
          def func(n):
              return np.identity(n)
          self.check_identity(func)
  
      def test_identity_dtype(self):
          for dtype in (np.complex64, np.int16, np.bool_, np.dtype('bool'),
                        'bool_'):
              def func(n):
                  return np.identity(n, dtype)
              self.check_identity(func)
  
      def test_like_dtype_non_const_str_kwarg(self):
  
          @njit
          def func(n, dt):
              return np.identity(n, dt)
  
          with self.assertRaises(TypingError) as raises:
              func(4, 'int32')
  
          excstr = str(raises.exception)
          self.assertIn('No match', excstr)
          regex = re.compile(r'\bidentity\(int.*?, unicode_type\)\B')
          self.assertRegex(excstr, regex)
  
  
  class TestNdEye(BaseTest):
  
      def test_eye_n(self):
          def func(n):
              return np.eye(n)
          self.check_outputs(func, [(1,), (3,)])
  
      def test_eye_n_dtype(self):
          # check None option, dtype class, instance of dtype class
          for dt in (None, np.complex128, np.complex64(1)):
              def func(n, dtype=dt):
                  return np.eye(n, dtype=dtype)
              self.check_outputs(func, [(1,), (3,)])
  
      def test_eye_n_m(self):
          def func(n, m):
              return np.eye(n, m)
          self.check_outputs(func, [(1, 2), (3, 2), (0, 3)])
  
      def check_eye_n_m_k(self, func):
          self.check_outputs(func, [(1, 2, 0),
                                    (3, 4, 1),
                                    (3, 4, -1),
                                    (4, 3, -2),
                                    (4, 3, -5),
                                    (4, 3, 5)])
  
      def test_eye_n_m_k(self):
          def func(n, m, k):
              return np.eye(n, m, k)
          self.check_eye_n_m_k(func)
  
      def test_eye_n_m_k_dtype(self):
          def func(n, m, k):
              return np.eye(N=n, M=m, k=k, dtype=np.int16)
          self.check_eye_n_m_k(func)
  
      def test_eye_n_m_k_dtype_instance(self):
          dtype = np.dtype('int16')
          def func(n, m, k):
              return np.eye(N=n, M=m, k=k, dtype=dtype)
          self.check_eye_n_m_k(func)
  
  
  class TestNdDiag(TestCase):
  
      def setUp(self):
          v = np.array([1, 2, 3])
          hv = np.array([[1, 2, 3]])
          vv = np.transpose(hv)
          self.vectors = [v, hv, vv]
          a3x4 = np.arange(12).reshape(3, 4)
          a4x3 = np.arange(12).reshape(4, 3)
          self.matricies = [a3x4, a4x3]
          def func(q):
              return np.diag(q)
          self.py = func
          self.jit = nrtjit(func)
  
          def func_kwarg(q, k=0):
              return np.diag(q, k=k)
          self.py_kw = func_kwarg
          self.jit_kw = nrtjit(func_kwarg)
  
      def check_diag(self, pyfunc, nrtfunc, *args, **kwargs):
          expected = pyfunc(*args, **kwargs)
          computed = nrtfunc(*args, **kwargs)
          self.assertEqual(computed.size, expected.size)
          self.assertEqual(computed.dtype, expected.dtype)
          # NOTE: stride not tested as np returns a RO view, nb returns new data
          np.testing.assert_equal(expected, computed)
  
      # create a diag matrix from a vector
      def test_diag_vect_create(self):
          for d in self.vectors:
              self.check_diag(self.py, self.jit, d)
  
      # create a diag matrix from a vector at a given offset
      def test_diag_vect_create_kwarg(self):
          for k in range(-10, 10):
              for d in self.vectors:
                  self.check_diag(self.py_kw, self.jit_kw, d, k=k)
  
      # extract the diagonal
      def test_diag_extract(self):
          for d in self.matricies:
              self.check_diag(self.py, self.jit, d)
  
      # extract a diagonal at a given offset
      def test_diag_extract_kwarg(self):
          for k in range(-4, 4):
              for d in self.matricies:
                  self.check_diag(self.py_kw, self.jit_kw, d, k=k)
  
      # check error handling
      def test_error_handling(self):
          d = np.array([[[1.]]])
          cfunc = nrtjit(self.py)
  
          # missing arg
          with self.assertRaises(TypeError):
              cfunc()
  
          # > 2d
          with self.assertRaises(TypingError):
              cfunc(d)
          with self.assertRaises(TypingError):
              dfunc = nrtjit(self.py_kw)
              dfunc(d, k=3)
  
  class TestNdArange(BaseTest):
  
      def test_linspace_2(self):
          def pyfunc(n, m):
              return np.linspace(n, m)
          self.check_outputs(pyfunc,
                             [(0, 4), (1, 100), (-3.5, 2.5), (-3j, 2+3j),
                              (2, 1), (1+0.5j, 1.5j)], exact=False)
  
      def test_linspace_3(self):
          def pyfunc(n, m, p):
              return np.linspace(n, m, p)
          self.check_outputs(pyfunc,
                             [(0, 4, 9), (1, 4, 3), (-3.5, 2.5, 8),
                              (-3j, 2+3j, 7), (2, 1, 0),
                              (1+0.5j, 1.5j, 5), (1, 1e100, 1)],
                             exact=False)
  
  
  class TestNpyEmptyKeyword(TestCase):
      def _test_with_dtype_kw(self, dtype):
          def pyfunc(shape):
              return np.empty(shape, dtype=dtype)
  
          shapes = [1, 5, 9]
  
          cfunc = nrtjit(pyfunc)
          for s in shapes:
              expected = pyfunc(s)
              got = cfunc(s)
              self.assertEqual(expected.dtype, got.dtype)
              self.assertEqual(expected.shape, got.shape)
  
      def test_with_dtype_kws(self):
          for dtype in [np.int32, np.float32, np.complex64, np.dtype('complex64')]:
              self._test_with_dtype_kw(dtype)
  
      def _test_with_shape_and_dtype_kw(self, dtype):
          def pyfunc(shape):
              return np.empty(shape=shape, dtype=dtype)
  
          shapes = [1, 5, 9]
  
          cfunc = nrtjit(pyfunc)
          for s in shapes:
              expected = pyfunc(s)
              got = cfunc(s)
              self.assertEqual(expected.dtype, got.dtype)
              self.assertEqual(expected.shape, got.shape)
  
      def test_with_shape_and_dtype_kws(self):
          for dtype in [np.int32, np.float32, np.complex64, np.dtype('complex64')]:
              self._test_with_shape_and_dtype_kw(dtype)
  
      def test_empty_no_args(self):
  
          def pyfunc():
              return np.empty()
  
          cfunc = nrtjit(pyfunc)
  
          # Trigger the compilation
          # That will cause a TypingError due to missing shape argument
          with self.assertRaises(TypingError):
              cfunc()
  
  
  class TestNpArray(MemoryLeakMixin, BaseTest):
  
      def test_0d(self):
          def pyfunc(arg):
              return np.array(arg)
  
          cfunc = nrtjit(pyfunc)
          got = cfunc(42)
          self.assertPreciseEqual(got, np.array(42, dtype=np.intp))
          got = cfunc(2.5)
          self.assertPreciseEqual(got, np.array(2.5))
  
      def test_0d_with_dtype(self):
          def pyfunc(arg):
              return np.array(arg, dtype=np.int16)
  
          self.check_outputs(pyfunc, [(42,), (3.5,)])
  
      def test_1d(self):
          def pyfunc(arg):
              return np.array(arg)
  
          cfunc = nrtjit(pyfunc)
          # A list
          got = cfunc([2, 3, 42])
          self.assertPreciseEqual(got, np.intp([2, 3, 42]))
          # A heterogeneous tuple
          got = cfunc((1.0, 2.5j, 42))
          self.assertPreciseEqual(got, np.array([1.0, 2.5j, 42]))
          # An empty tuple
          got = cfunc(())
          self.assertPreciseEqual(got, np.float64(()))
  
      def test_1d_with_dtype(self):
          def pyfunc(arg):
              return np.array(arg, dtype=np.float32)
  
          self.check_outputs(pyfunc,
                             [([2, 42],),
                              ([3.5, 1.0],),
                              ((1, 3.5, 42),),
                              ((),),
                              ])
  
      def test_1d_with_str_dtype(self):
          def pyfunc(arg):
              return np.array(arg, dtype='float32')
  
          self.check_outputs(pyfunc,
                             [([2, 42],),
                              ([3.5, 1.0],),
                              ((1, 3.5, 42),),
                              ((),),
                              ])
  
      def test_1d_with_non_const_str_dtype(self):
  
          @njit
          def func(arg, dt):
              return np.array(arg, dtype=dt)
  
          with self.assertRaises(TypingError) as raises:
              func((5, 3), 'int32')
  
          excstr = str(raises.exception)
          self.assertIn('No match', excstr)
          restr = r'\barray\(UniTuple\(int.*? x 2\), dtype=unicode_type\)\B'
          regex = re.compile(restr)
          self.assertRegex(excstr, regex)
  
      def test_2d(self):
          def pyfunc(arg):
              return np.array(arg)
  
          cfunc = nrtjit(pyfunc)
          # A list of tuples
          got = cfunc([(1, 2), (3, 4)])
          self.assertPreciseEqual(got, np.intp([[1, 2], [3, 4]]))
          got = cfunc([(1, 2.5), (3, 4.5)])
          self.assertPreciseEqual(got, np.float64([[1, 2.5], [3, 4.5]]))
          # A tuple of lists
          got = cfunc(([1, 2], [3, 4]))
          self.assertPreciseEqual(got, np.intp([[1, 2], [3, 4]]))
          got = cfunc(([1, 2], [3.5, 4.5]))
          self.assertPreciseEqual(got, np.float64([[1, 2], [3.5, 4.5]]))
          # A tuple of tuples
          got = cfunc(((1.5, 2), (3.5, 4.5)))
          self.assertPreciseEqual(got, np.float64([[1.5, 2], [3.5, 4.5]]))
          got = cfunc(((), ()))
          self.assertPreciseEqual(got, np.float64(((), ())))
  
      def test_2d_with_dtype(self):
          def pyfunc(arg):
              return np.array(arg, dtype=np.int32)
  
          cfunc = nrtjit(pyfunc)
          got = cfunc([(1, 2.5), (3, 4.5)])
          self.assertPreciseEqual(got, np.int32([[1, 2], [3, 4]]))
  
      def test_raises(self):
  
          def pyfunc(arg):
              return np.array(arg)
  
          cfunc = nrtjit(pyfunc)
  
          @contextlib.contextmanager
          def check_raises(msg):
              with self.assertRaises(TypingError) as raises:
                  yield
              self.assertIn(msg, str(raises.exception))
  
          with check_raises(('array(float64, 1d, C) not allowed in a '
                             'homogeneous sequence')):
              cfunc(np.array([1.]))
  
          with check_raises(('type Tuple(int64, reflected list(int64)<iv=None>) '
                            'does not have a regular shape')):
              cfunc((np.int64(1), [np.int64(2)]))
  
          with check_raises(
                  "cannot convert Tuple(int64, Record(a[type=int32;offset=0],"
                  "b[type=float32;offset=4];8;False)) to a homogeneous type",
              ):
              st = np.dtype([('a', 'i4'), ('b', 'f4')])
              val = np.zeros(1, dtype=st)[0]
              cfunc(((1, 2), (np.int64(1), val)))
  
  
  class TestNpConcatenate(MemoryLeakMixin, TestCase):
      """
      Tests for np.concatenate().
      """
  
      def _3d_arrays(self):
          a = np.arange(24).reshape((4, 3, 2))
          b = a + 10
          c = (b + 10).copy(order='F')
          d = (c + 10)[::-1]
          e = (d + 10)[...,::-1]
          return a, b, c, d, e
  
      @contextlib.contextmanager
      def assert_invalid_sizes_over_dim(self, axis):
          with self.assertRaises(ValueError) as raises:
              yield
          self.assertIn("input sizes over dimension %d do not match" % axis,
                        str(raises.exception))
  
      def test_3d(self):
          pyfunc = np_concatenate2
          cfunc = nrtjit(pyfunc)
  
          def check(a, b, c, axis):
              for ax in (axis, -3 + axis):
                  expected = pyfunc(a, b, c, axis=ax)
                  got = cfunc(a, b, c, axis=ax)
                  self.assertPreciseEqual(got, expected)
  
          def check_all_axes(a, b, c):
              for axis in range(3):
                  check(a, b, c, axis)
  
          a, b, c, d, e = self._3d_arrays()
  
          # Inputs with equal sizes
          # C, C, C
          check_all_axes(a, b, b)
          # C, C, F
          check_all_axes(a, b, c)
          # F, F, F
          check_all_axes(a.T, b.T, a.T)
          # F, F, C
          check_all_axes(a.T, b.T, c.T)
          # F, F, A
          check_all_axes(a.T, b.T, d.T)
          # A, A, A
          # (note Numpy may select the layout differently for other inputs)
          check_all_axes(d.T, e.T, d.T)
  
          # Inputs with compatible sizes
          check(a[1:], b, c[::-1], axis=0)
          check(a, b[:,1:], c, axis=1)
          check(a, b, c[:,:,1:], axis=2)
  
          # Different but compatible dtypes
          check_all_axes(a, b.astype(np.float64), b)
  
          # Exceptions leak references
          self.disable_leak_check()
  
          # Incompatible sizes
          for axis in (1, 2, -2, -1):
              with self.assert_invalid_sizes_over_dim(0):
                  cfunc(a[1:], b, b, axis)
          for axis in (0, 2, -3, -1):
              with self.assert_invalid_sizes_over_dim(1):
                  cfunc(a, b[:,1:], b, axis)
  
      def test_3d_no_axis(self):
          pyfunc = np_concatenate1
          cfunc = nrtjit(pyfunc)
  
          def check(a, b, c):
              expected = pyfunc(a, b, c)
              got = cfunc(a, b, c)
              self.assertPreciseEqual(got, expected)
  
          a, b, c, d, e = self._3d_arrays()
  
          # Inputs with equal sizes
          # C, C, C
          check(a, b, b)
          # C, C, F
          check(a, b, c)
          # F, F, F
          check(a.T, b.T, a.T)
          # F, F, C
          check(a.T, b.T, c.T)
          # F, F, A
          check(a.T, b.T, d.T)
          # A, A, A
          # (note Numpy may select the layout differently for other inputs)
          check(d.T, e.T, d.T)
  
          # Inputs with compatible sizes
          check(a[1:], b, c[::-1])
  
          # Exceptions leak references
          self.disable_leak_check()
  
          # Incompatible sizes
          with self.assert_invalid_sizes_over_dim(1):
              cfunc(a, b[:,1:], b)
  
      def test_typing_errors(self):
          pyfunc = np_concatenate1
          cfunc = nrtjit(pyfunc)
  
          a = np.arange(15)
          b = a.reshape((3, 5))
          c = a.astype(np.dtype([('x', np.int8)]))
          d = np.array(42)
  
          # Different dimensionalities
          with self.assertTypingError() as raises:
              cfunc(a, b, b)
          self.assertIn("all the input arrays must have same number of dimensions",
                        str(raises.exception))
  
          # Incompatible dtypes
          with self.assertTypingError() as raises:
              cfunc(a, c, c)
          self.assertIn("input arrays must have compatible dtypes",
                        str(raises.exception))
  
          # 0-d arrays
          with self.assertTypingError() as raises:
              cfunc(d, d, d)
          self.assertIn("zero-dimensional arrays cannot be concatenated",
                        str(raises.exception))
  
  
  @unittest.skipUnless(hasattr(np, "stack"), "this Numpy doesn't have np.stack()")
  class TestNpStack(MemoryLeakMixin, TestCase):
      """
      Tests for np.stack().
      """
  
      def _3d_arrays(self):
          a = np.arange(24).reshape((4, 3, 2))
          b = a + 10
          c = (b + 10).copy(order='F')
          d = (c + 10)[::-1]
          e = (d + 10)[...,::-1]
          return a, b, c, d, e
  
      @contextlib.contextmanager
      def assert_invalid_sizes(self):
          with self.assertRaises(ValueError) as raises:
              yield
          self.assertIn("all input arrays must have the same shape",
                        str(raises.exception))
  
      def check_stack(self, pyfunc, cfunc, args):
          expected = pyfunc(*args)
          got = cfunc(*args)
          # Numba doesn't choose the same layout as Numpy.
          # We would like to check the result is contiguous, but we can't
          # rely on the "flags" attribute when there are 1-sized
          # dimensions.
          self.assertEqual(got.shape, expected.shape)
          self.assertPreciseEqual(got.flatten(), expected.flatten())
  
      def check_3d(self, pyfunc, cfunc, generate_starargs):
          def check(a, b, c, args):
              self.check_stack(pyfunc, cfunc, (a, b, c) + args)
  
          def check_all_axes(a, b, c):
              for args in generate_starargs():
                  check(a, b, c, args)
  
          a, b, c, d, e = self._3d_arrays()
  
          # C, C, C
          check_all_axes(a, b, b)
          # C, C, F
          check_all_axes(a, b, c)
          # F, F, F
          check_all_axes(a.T, b.T, a.T)
          # F, F, C
          check_all_axes(a.T, b.T, c.T)
          # F, F, A
          check_all_axes(a.T, b.T, d.T)
          # A, A, A
          check_all_axes(d.T, e.T, d.T)
  
          # Different but compatible dtypes
          check_all_axes(a, b.astype(np.float64), b)
  
      def check_runtime_errors(self, cfunc, generate_starargs):
          # Exceptions leak references
          self.assert_no_memory_leak()
          self.disable_leak_check()
  
          # Inputs have different shapes
          a, b, c, d, e = self._3d_arrays()
          with self.assert_invalid_sizes():
              args = next(generate_starargs())
              cfunc(a[:-1], b, c, *args)
  
      def test_3d(self):
          """
          stack(3d arrays, axis)
          """
          pyfunc = np_stack2
          cfunc = nrtjit(pyfunc)
  
          def generate_starargs():
              for axis in range(3):
                  yield (axis,)
                  yield (-3 + axis,)
  
          self.check_3d(pyfunc, cfunc, generate_starargs)
          self.check_runtime_errors(cfunc, generate_starargs)
  
      def test_3d_no_axis(self):
          """
          stack(3d arrays)
          """
          pyfunc = np_stack1
          cfunc = nrtjit(pyfunc)
  
          def generate_starargs():
              yield()
  
          self.check_3d(pyfunc, cfunc, generate_starargs)
          self.check_runtime_errors(cfunc, generate_starargs)
  
      def test_0d(self):
          """
          stack(0d arrays)
          """
          pyfunc = np_stack1
          cfunc = nrtjit(pyfunc)
  
          a = np.array(42)
          b = np.array(-5j)
          c = np.array(True)
  
          self.check_stack(pyfunc, cfunc, (a, b, c))
  
      def check_xxstack(self, pyfunc, cfunc):
          """
          3d and 0d tests for hstack(), vstack(), dstack().
          """
          def generate_starargs():
              yield()
  
          self.check_3d(pyfunc, cfunc, generate_starargs)
          # 0d
          a = np.array(42)
          b = np.array(-5j)
          c = np.array(True)
          self.check_stack(pyfunc, cfunc, (a, b, a))
  
      def test_hstack(self):
          pyfunc = np_hstack
          cfunc = nrtjit(pyfunc)
  
          self.check_xxstack(pyfunc, cfunc)
          # 1d
          a = np.arange(5)
          b = np.arange(6) + 10
          self.check_stack(pyfunc, cfunc, (a, b, b))
          # 2d
          a = np.arange(6).reshape((2, 3))
          b = np.arange(8).reshape((2, 4)) + 100
          self.check_stack(pyfunc, cfunc, (a, b, a))
  
      def test_vstack(self):
          pyfunc = np_vstack
          cfunc = nrtjit(pyfunc)
  
          self.check_xxstack(pyfunc, cfunc)
          # 1d
          a = np.arange(5)
          b = a + 10
          self.check_stack(pyfunc, cfunc, (a, b, b))
          # 2d
          a = np.arange(6).reshape((3, 2))
          b = np.arange(8).reshape((4, 2)) + 100
          self.check_stack(pyfunc, cfunc, (a, b, b))
  
      def test_dstack(self):
          pyfunc = np_dstack
          cfunc = nrtjit(pyfunc)
  
          self.check_xxstack(pyfunc, cfunc)
          # 1d
          a = np.arange(5)
          b = a + 10
          self.check_stack(pyfunc, cfunc, (a, b, b))
          # 2d
          a = np.arange(12).reshape((3, 4))
          b = a + 100
          self.check_stack(pyfunc, cfunc, (a, b, b))
  
      def test_column_stack(self):
          pyfunc = np_column_stack
          cfunc = nrtjit(pyfunc)
  
          a = np.arange(4)
          b = a + 10
          c = np.arange(12).reshape((4, 3))
          self.check_stack(pyfunc, cfunc, (a, b, c))
  
          # Exceptions leak references
          self.assert_no_memory_leak()
          self.disable_leak_check()
  
          # Invalid dims
          a = np.array(42)
          with self.assertTypingError():
              cfunc((a, a, a))
          a = a.reshape((1, 1, 1))
          with self.assertTypingError():
              cfunc((a, a, a))
  
  
  def benchmark_refct_speed():
      def pyfunc(x, y, t):
          """Swap array x and y for t number of times
          """
          for i in range(t):
              x, y = y, x
          return x, y
  
      cfunc = nrtjit(pyfunc)
  
      x = np.random.random(100)
      y = np.random.random(100)
      t = 10000
  
      def bench_pyfunc():
          pyfunc(x, y, t)
  
      def bench_cfunc():
          cfunc(x, y, t)
  
      python_time = utils.benchmark(bench_pyfunc)
      numba_time = utils.benchmark(bench_cfunc)
      print(python_time)
      print(numba_time)
  
  
  if __name__ == "__main__":
      unittest.main()
  