feat: add N-D concatenate with axis support — bit-exact with numpy

Add axis-aware N-D concatenate to the native core (core.h), updating
the pybind11 wrapper (core_py.h) and bindings (module.cpp) to accept
an axis parameter.

Native implementation uses leading-slice block copies — each slice
contributes contiguous elements per array, so a single memcpy per
array per slice suffices. Per-array strides correctly account for
differing axis dimension sizes.

Also improve vstack/hstack wrappers:
- vstack: 1D arrays reshaped to (1,N) before stacking
- hstack: uses axis=1 for 2D+ arrays

Adds 30+ new concatenate tests covering 1D–5D, all axes, float32/64,
large arrays, identity, zeros/ones, and edge cases.

Author: peng.li24

numpy/core.h

@@ -532,6 +532,7 @@ inline void stack(const T* const* arrays, T* dst, size_t n_arrays, size_t elem_s
 }
 
 /// numpy.concatenate((a1, a2, ...), axis=0, out=None, dtype=None, casting=...)
+/// 1D flat: arrays are treated as flat buffers, concatenated sequentially.
 template<typename T>
 inline void concatenate(const T* const* arrays, T* dst, const size_t* sizes, size_t n_arrays) {
     size_t off = 0;
@@ -541,6 +542,97 @@ inline void concatenate(const T* const* arrays, T* dst, const size_t* sizes, siz
     }
 }
 
+/// numpy.concatenate((a1, a2, ...), axis=0, ...) — N-D with axis support.
+/// All arrays must have identical shape except along `axis`.
+/// `shape` is the representative common shape (use first array's shape);
+/// `axis_sizes[i]` gives the size of array i along the concatenation axis.
+///
+/// Strategy: iterate over "leading slices" (product of dims before axis).
+/// Within each slice, every array contributes a contiguous block of
+/// `axis_sizes[i] * trailing` elements.  Since the elements are C-contiguous
+/// within each slice, a single memcpy per array per slice suffices.
+template<typename T>
+inline void concatenate(const T* const* arrays, T* dst,
+                         const ptrdiff_t* shape, int ndim, int axis,
+                         const size_t* axis_sizes, size_t n_arrays) {
+    if (n_arrays == 0 || ndim == 0) return;
+
+    // Normalize axis
+    if (axis < 0) axis += ndim;
+
+    // Trailing product = product of dims after axis (also = stride along axis)
+    ptrdiff_t trailing = 1;
+    for (int d = axis + 1; d < ndim; ++d) trailing *= shape[d];
+
+    // Total output axis size
+    ptrdiff_t out_axis = 0;
+    for (size_t i = 0; i < n_arrays; ++i)
+        out_axis += static_cast<ptrdiff_t>(axis_sizes[i]);
+
+    // Per-array full strides (differ because axis dim sizes differ)
+    // C-contiguous: stride[d] = stride[d+1] * size_of_dim[d+1].
+    // Use axis_sizes[k] when d+1 == axis, common shape otherwise.
+    std::vector<std::vector<ptrdiff_t>> in_stride(n_arrays);
+    for (size_t k = 0; k < n_arrays; ++k) {
+        in_stride[k].resize(ndim);
+        in_stride[k][ndim - 1] = 1;
+        for (int d = ndim - 2; d >= 0; --d) {
+            ptrdiff_t s = (d + 1 == axis)
+                ? static_cast<ptrdiff_t>(axis_sizes[k])
+                : shape[d + 1];
+            in_stride[k][d] = in_stride[k][d + 1] * s;
+        }
+    }
+
+    // Output strides
+    std::vector<ptrdiff_t> out_shape(shape, shape + ndim);
+    out_shape[axis] = out_axis;
+    std::vector<ptrdiff_t> out_stride(ndim);
+    out_stride[ndim - 1] = 1;
+    for (int d = ndim - 2; d >= 0; --d)
+        out_stride[d] = out_stride[d + 1] * out_shape[d + 1];
+
+    // Per-array per-slice element count (contiguous elements contributed per slice)
+    std::vector<size_t> slice_n(n_arrays);
+    for (size_t i = 0; i < n_arrays; ++i)
+        slice_n[i] = static_cast<size_t>(axis_sizes[i]) * static_cast<size_t>(trailing);
+
+    // Number of leading slices
+    ptrdiff_t n_slices = 1;
+    for (int d = 0; d < axis; ++d) n_slices *= shape[d];
+
+    // Total per-array byte size of one slice (for output position stepping)
+    ptrdiff_t out_slice_bytes = static_cast<ptrdiff_t>(
+        static_cast<size_t>(out_axis) * static_cast<size_t>(trailing) * sizeof(T));
+
+    // For each leading slice, copy contiguous blocks from each array
+    for (ptrdiff_t s = 0; s < n_slices; ++s) {
+        // Decompose slice index → multi-index for dims 0..axis-1
+        ptrdiff_t rem = s;
+
+        // Per-array leading offset within the array
+        std::vector<size_t> in_off(n_arrays, 0);
+
+        for (int d = axis - 1; d >= 0; --d) {
+            ptrdiff_t idx = rem % shape[d];
+            rem /= shape[d];
+            for (size_t k = 0; k < n_arrays; ++k)
+                in_off[k] += static_cast<size_t>(idx) * static_cast<size_t>(in_stride[k][d]);
+        }
+
+        // Output position for this slice
+        char* out_slice_start = reinterpret_cast<char*>(dst) + s * out_slice_bytes;
+        size_t out_byte_off = 0;
+
+        for (size_t i = 0; i < n_arrays; ++i) {
+            size_t bytes = slice_n[i] * sizeof(T);
+            std::memcpy(out_slice_start + out_byte_off,
+                        arrays[i] + in_off[i], bytes);
+            out_byte_off += bytes;
+        }
+    }
+}
+
 /// numpy.where(condition, x, y) — scalar x, y
 template<typename T>
 inline void where_scalar(const bool* cond, T* dst, size_t n, T x, T y) {

pycpp/core_py.h

@@ -884,28 +884,93 @@ py::array_t<T> stack(const std::vector<py::array_t<T>>& arrays) {
 
 /// numpy.concatenate((a1, a2, ...), axis=0, out=None, dtype=None, casting=...)
 template<typename T>
-py::array_t<T> concatenate(const std::vector<py::array_t<T>>& arrays) {
+py::array_t<T> concatenate(const std::vector<py::array_t<T>>& arrays, int axis = 0) {
     if (arrays.empty()) return py::array_t<T>{};
-    py::ssize_t total = 0;
-    for (const auto& arr : arrays) total += arr.request().size;
-    py::array_t<T> result({total});
-    T* dst = static_cast<T*>(result.request().ptr);
-    py::ssize_t off = 0;
+
+    auto buf0 = arrays[0].request();
+    int ndim = static_cast<int>(buf0.ndim);
+
+    if (axis < 0) axis += ndim;
+    if (axis < 0 || axis >= ndim)
+        throw std::invalid_argument("concatenate: axis out of range");
+
+    // Validate that all arrays have same number of dimensions
     for (const auto& arr : arrays) {
-        auto buf = arr.request();
-        std::memcpy(dst + off, static_cast<const T*>(buf.ptr), buf.size * sizeof(T));
-        off += buf.size;
+        if (arr.request().ndim != ndim)
+            throw std::invalid_argument("concatenate: all arrays must have same number of dimensions");
+    }
+
+    // Collect shape (from first array) and per-array axis sizes
+    std::vector<ptrdiff_t> shape(ndim);
+    for (int d = 0; d < ndim; ++d) shape[d] = buf0.shape[d];
+
+    std::vector<size_t> axis_sizes(arrays.size());
+    for (size_t i = 0; i < arrays.size(); ++i) {
+        auto buf = arrays[i].request();
+        axis_sizes[i] = static_cast<size_t>(buf.shape[axis]);
+    }
+
+    // Validate non-axis dimensions match
+    for (size_t i = 0; i < arrays.size(); ++i) {
+        auto buf = arrays[i].request();
+        for (int d = 0; d < ndim; ++d) {
+            if (d == axis) continue;
+            if (buf.shape[d] != shape[d])
+                throw std::invalid_argument(
+                    "concatenate: all arrays must have same shape except along axis");
+        }
     }
+
+    // Compute output shape
+    std::vector<ptrdiff_t> out_shape = shape;
+    ptrdiff_t total_axis = 0;
+    for (size_t i = 0; i < arrays.size(); ++i)
+        total_axis += static_cast<ptrdiff_t>(axis_sizes[i]);
+    out_shape[axis] = total_axis;
+
+    std::vector<py::ssize_t> py_out_shape(out_shape.begin(), out_shape.end());
+    py::array_t<T> result(py_out_shape);
+    T* dst = static_cast<T*>(result.request().ptr);
+
+    // Build pointer array
+    std::vector<const T*> ptrs(arrays.size());
+    for (size_t i = 0; i < arrays.size(); ++i)
+        ptrs[i] = static_cast<const T*>(arrays[i].request().ptr);
+
+    numpy::concatenate(ptrs.data(), dst, shape.data(), ndim, axis,
+                       axis_sizes.data(), arrays.size());
     return result;
 }
 
 /// numpy.vstack(tup, *, dtype=None, casting=...)
 template<typename T>
-py::array_t<T> vstack(const std::vector<py::array_t<T>>& arrays) { return stack(arrays); }
+py::array_t<T> vstack(const std::vector<py::array_t<T>>& arrays) {
+    if (arrays.empty()) return py::array_t<T>{};
+    int ndim = static_cast<int>(arrays[0].request().ndim);
+    if (ndim == 1) {
+        // numpy.vstack: 1D arrays are reshaped to (1, N) before stacking
+        auto buf0 = arrays[0].request();
+        py::array_t<T> result({static_cast<py::ssize_t>(arrays.size()), static_cast<py::ssize_t>(buf0.size)});
+        T* dst = static_cast<T*>(result.request().ptr);
+        for (size_t i = 0; i < arrays.size(); ++i) {
+            auto buf = arrays[i].request();
+            std::memcpy(dst + i * buf0.size, static_cast<const T*>(buf.ptr),
+                        buf.size * sizeof(T));
+        }
+        return result;
+    }
+    return concatenate(arrays, 0);
+}
 
 /// numpy.hstack(tup, *, dtype=None, casting=...)
 template<typename T>
-py::array_t<T> hstack(const std::vector<py::array_t<T>>& arrays) { return concatenate(arrays); }
+py::array_t<T> hstack(const std::vector<py::array_t<T>>& arrays) {
+    if (arrays.empty()) return py::array_t<T>{};
+    int ndim = static_cast<int>(arrays[0].request().ndim);
+    // 1D arrays: hstack is identical to concatenate along axis=0
+    // 2D+ arrays: hstack concatenates along axis=1
+    return concatenate(arrays, (ndim == 1) ? 0 : 1);
+}
 
 /// numpy.where(condition, x, y) — scalar x, y
 template<typename T>

tests/module.cpp

@@ -196,7 +196,12 @@ PYBIND11_MODULE(numpycpp, m) {
           py::arg("arr"), py::arg("n") = 1, py::arg("axis") = -1);
     m.def("diff", static_cast<py::array_t<float>(*)(const py::array_t<float>&, int, int)>(&numpy::diff),
           py::arg("arr"), py::arg("n") = 1, py::arg("axis") = -1);
-    BIND_F_STACK(stack); BIND_F_STACK(concatenate); BIND_F_STACK(vstack); BIND_F_STACK(hstack);
+    BIND_F_STACK(stack); BIND_F_STACK(vstack); BIND_F_STACK(hstack);
+
+    m.def("concatenate", static_cast<py::array_t<double>(*)(const std::vector<py::array_t<double>>&, int)>(&numpy::concatenate),
+          py::arg("arrays"), py::arg("axis") = 0);
+    m.def("concatenate", static_cast<py::array_t<float>(*)(const std::vector<py::array_t<float>>&, int)>(&numpy::concatenate),
+          py::arg("arrays"), py::arg("axis") = 0);
 
     m.def("where", static_cast<py::array_t<double>(*)(const py::array_t<bool>&, double, double)>(&numpy::where));
     m.def("where", static_cast<py::array_t<float>(*)(const py::array_t<bool>&, float, float)>(&numpy::where));

tests/test_all.py

@@ -580,17 +580,136 @@ def test_stack(cpp, dtype):
     arrays = [random_array((3,), seed=i, dtype=dtype) for i in range(4)]
     assert_bit_aligned(cpp.stack(arrays), np.stack(arrays), "stack")
 
-def test_concatenate(cpp, dtype):
+def test_concatenate_1d(cpp, dtype):
     arrays = [random_array((3,), seed=i, dtype=dtype) for i in range(3)]
-    assert_bit_aligned(cpp.concatenate(arrays), np.concatenate(arrays), "concatenate")
+    assert_bit_aligned(cpp.concatenate(arrays), np.concatenate(arrays), "concatenate 1d")
+
+def test_concatenate_2d_axis0(cpp, dtype):
+    arrays = [random_array((2, 3), seed=i, dtype=dtype) for i in range(3)]
+    assert_bit_aligned(cpp.concatenate(arrays, 0), np.concatenate(arrays, axis=0), "concatenate 2d axis=0")
+    # Verify default axis=0
+    assert_bit_aligned(cpp.concatenate(arrays), np.concatenate(arrays), "concatenate 2d default axis")
+
+def test_concatenate_2d_axis1(cpp, dtype):
+    arrays = [random_array((3, 2), seed=i, dtype=dtype) for i in range(3)]
+    assert_bit_aligned(cpp.concatenate(arrays, 1), np.concatenate(arrays, axis=1), "concatenate 2d axis=1")
+
+def test_concatenate_2d_axis_neg1(cpp, dtype):
+    arrays = [random_array((3, 2), seed=i, dtype=dtype) for i in range(3)]
+    assert_bit_aligned(cpp.concatenate(arrays, -1), np.concatenate(arrays, axis=-1), "concatenate 2d axis=-1")
+
+def test_concatenate_3d_axis0(cpp, dtype):
+    arrays = [random_array((2, 3, 4), seed=i, dtype=dtype) for i in range(2)]
+    assert_bit_aligned(cpp.concatenate(arrays, 0), np.concatenate(arrays, axis=0), "concatenate 3d axis=0")
+
+def test_concatenate_3d_axis1(cpp, dtype):
+    arrays = [random_array((3, 2, 4), seed=i, dtype=dtype) for i in range(2)]
+    assert_bit_aligned(cpp.concatenate(arrays, 1), np.concatenate(arrays, axis=1), "concatenate 3d axis=1")
+
+def test_concatenate_3d_axis2(cpp, dtype):
+    arrays = [random_array((3, 4, 2), seed=i, dtype=dtype) for i in range(2)]
+    assert_bit_aligned(cpp.concatenate(arrays, 2), np.concatenate(arrays, axis=2), "concatenate 3d axis=2")
+
+def test_concatenate_two_arrays(cpp, dtype):
+    arrays = [random_array((5,), seed=0, dtype=dtype), random_array((7,), seed=1, dtype=dtype)]
+    assert_bit_aligned(cpp.concatenate(arrays), np.concatenate(arrays), "concatenate two")
+
+def test_concatenate_single(cpp, dtype):
+    arrays = [random_array((5,), dtype=dtype)]
+    assert_bit_aligned(cpp.concatenate(arrays), np.concatenate(arrays), "concatenate single")
 
 def test_vstack(cpp, dtype):
     arrays = [random_array((1, 3), seed=i, dtype=dtype) for i in range(4)]
     assert_bit_aligned(cpp.vstack(arrays), np.vstack(arrays), "vstack")
 
+def test_vstack_1d(cpp, dtype):
+    arrays = [random_array((3,), seed=i, dtype=dtype) for i in range(4)]
+    assert_bit_aligned(cpp.vstack(arrays), np.vstack(arrays), "vstack 1d")
+
 def test_hstack(cpp, dtype):
     arrays = [random_array((3,), seed=i, dtype=dtype) for i in range(3)]
-    assert_bit_aligned(cpp.hstack(arrays), np.hstack(arrays), "hstack")
+    assert_bit_aligned(cpp.hstack(arrays), np.hstack(arrays), "hstack 1d")
+
+def test_hstack_2d(cpp, dtype):
+    arrays = [random_array((3, 2), seed=i, dtype=dtype) for i in range(3)]
+    assert_bit_aligned(cpp.hstack(arrays), np.hstack(arrays), "hstack 2d")
+
+# -- Concatenate complex / edge-case tests ----------------------------------
+
+def test_concatenate_4d_axis0(cpp, dtype):
+    arrays = [random_array((2, 3, 4, 5), seed=i, dtype=dtype) for i in range(2)]
+    assert_bit_aligned(cpp.concatenate(arrays, 0), np.concatenate(arrays, axis=0), "concatenate 4d axis=0")
+
+def test_concatenate_4d_axis2(cpp, dtype):
+    arrays = [random_array((2, 3, 2, 5), seed=i, dtype=dtype) for i in range(2)]
+    assert_bit_aligned(cpp.concatenate(arrays, 2), np.concatenate(arrays, axis=2), "concatenate 4d axis=2")
+
+def test_concatenate_4d_axis_neg2(cpp, dtype):
+    arrays = [random_array((2, 3, 2, 5), seed=i, dtype=dtype) for i in range(2)]
+    assert_bit_aligned(cpp.concatenate(arrays, -2), np.concatenate(arrays, axis=-2), "concatenate 4d axis=-2")
+
+def test_concatenate_unequal_axis_sizes(cpp, dtype):
+    """Concatenate arrays of different sizes along the concatenation axis."""
+    a = random_array((3, 2), seed=1, dtype=dtype)
+    b = random_array((3, 4), seed=2, dtype=dtype)
+    c = random_array((3, 1), seed=3, dtype=dtype)
+    assert_bit_aligned(cpp.concatenate([a, b, c], 1),
+                       np.concatenate([a, b, c], axis=1), "concat unequal axis sizes")
+
+def test_concatenate_many_arrays(cpp, dtype):
+    """Concatenate 10 arrays along axis=0."""
+    arrays = [random_array((3,), seed=i, dtype=dtype) for i in range(10)]
+    assert_bit_aligned(cpp.concatenate(arrays), np.concatenate(arrays), "concat 10 arrays")
+
+def test_concatenate_large_3d(cpp, dtype):
+    """Large 3D concatenation along middle axis."""
+    arrays = [random_array((50, 20, 30), seed=i, dtype=dtype) for i in range(3)]
+    assert_bit_aligned(cpp.concatenate(arrays, 1), np.concatenate(arrays, axis=1), "concat large 3d axis=1")
+
+def test_concatenate_large_2d_axis0(cpp, dtype):
+    """Large 2D concatenation — 500 rows each, 4 arrays."""
+    arrays = [random_array((500, 10), seed=i, dtype=dtype) for i in range(4)]
+    assert_bit_aligned(cpp.concatenate(arrays, 0), np.concatenate(arrays, axis=0), "concat large 2d axis=0")
+
+def test_concatenate_large_2d_axis1(cpp, dtype):
+    """Large 2D concatenation — 500 cols each, 3 arrays."""
+    arrays = [random_array((10, 500), seed=i, dtype=dtype) for i in range(3)]
+    assert_bit_aligned(cpp.concatenate(arrays, 1), np.concatenate(arrays, axis=1), "concat large 2d axis=1")
+
+def test_concatenate_identity(cpp, dtype):
+    """Concatenating a single array returns identical copy."""
+    a = random_array((3, 4), seed=42, dtype=dtype)
+    assert_bit_aligned(cpp.concatenate([a], 0), np.concatenate([a], axis=0), "concat identity")
+    assert_bit_aligned(cpp.concatenate([a], 1), np.concatenate([a], axis=1), "concat identity axis=1")
+
+def test_concatenate_zeros(cpp, dtype):
+    """Concatenate arrays of zeros."""
+    a = np.zeros((2, 3), dtype=dtype)
+    b = np.zeros((2, 5), dtype=dtype)
+    assert_bit_aligned(cpp.concatenate([a, b], 1), np.concatenate([a, b], axis=1), "concat zeros")
+
+def test_concatenate_ones(cpp, dtype):
+    """Concatenate arrays of ones."""
+    a = np.ones((3, 2), dtype=dtype)
+    b = np.ones((5, 2), dtype=dtype)
+    assert_bit_aligned(cpp.concatenate([a, b], 0), np.concatenate([a, b], axis=0), "concat ones")
+
+def test_concatenate_3d_axis_neg2(cpp, dtype):
+    """3D concatenate along axis=-2 (middle axis)."""
+    arrays = [random_array((2, 3, 4), seed=i, dtype=dtype) for i in range(3)]
+    assert_bit_aligned(cpp.concatenate(arrays, -2), np.concatenate(arrays, axis=-2), "concat 3d axis=-2")
+
+def test_concatenate_3d_axis_neg3(cpp, dtype):
+    """3D concatenate along axis=-3 (first axis)."""
+    arrays = [random_array((2, 3, 4), seed=i, dtype=dtype) for i in range(2)]
+    assert_bit_aligned(cpp.concatenate(arrays, -3), np.concatenate(arrays, axis=-3), "concat 3d axis=-3")
+
+def test_concatenate_5d(cpp, dtype):
+    """5D concatenate along various axes."""
+    arrays = [random_array((2, 3, 2, 3, 2), seed=i, dtype=dtype) for i in range(2)]
+    assert_bit_aligned(cpp.concatenate(arrays, 0), np.concatenate(arrays, axis=0), "concat 5d axis=0")
+    assert_bit_aligned(cpp.concatenate(arrays, 2), np.concatenate(arrays, axis=2), "concat 5d axis=2")
+    assert_bit_aligned(cpp.concatenate(arrays, -1), np.concatenate(arrays, axis=-1), "concat 5d axis=-1")
 
 def test_where_scalar(cpp, dtype):
     cond = np.array([True, False, True, False, True])