from __future__ import annotations

from ._types import Optional, Tuple, Union, array

import numpy as np

def concat(arrays: Tuple[array], /, *, axis: Optional[int] = 0) -> array:
    """
    Array API compatible wrapper for :py:func:`np.concatenate <numpy.concatenate>`.

    See its docstring for more information.
    """
    # Note: the function name is different here
    return np.concatenate(arrays, axis=axis)

def expand_dims(x: array, axis: int, /) -> array:
    """
    Array API compatible wrapper for :py:func:`np.expand_dims <numpy.expand_dims>`.

    See its docstring for more information.
    """
    return np.expand_dims._implementation(x, axis)

def flip(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> array:
    """
    Array API compatible wrapper for :py:func:`np.flip <numpy.flip>`.

    See its docstring for more information.
    """
    return np.flip._implementation(x, axis=axis)

def reshape(x: array, shape: Tuple[int, ...], /) -> array:
    """
    Array API compatible wrapper for :py:func:`np.reshape <numpy.reshape>`.

    See its docstring for more information.
    """
    return np.reshape._implementation(x, shape)

def roll(x: array, shift: Union[int, Tuple[int, ...]], /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> array:
    """
    Array API compatible wrapper for :py:func:`np.roll <numpy.roll>`.

    See its docstring for more information.
    """
    return np.roll._implementation(x, shift, axis=axis)

def squeeze(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> array:
    """
    Array API compatible wrapper for :py:func:`np.squeeze <numpy.squeeze>`.

    See its docstring for more information.
    """
    return np.squeeze._implementation(x, axis=axis)

def stack(arrays: Tuple[array], /, *, axis: int = 0) -> array:
    """
    Array API compatible wrapper for :py:func:`np.stack <numpy.stack>`.

    See its docstring for more information.
    """
    return np.stack._implementation(arrays, axis=axis)

@@ -78,16 +78,51 @@
 n, array
+%0Afrom ._array_object import ndarray
 %0A%0Aimport
@@ -350,24 +350,69 @@
 on.%0A    %22%22%22%0A
+    arrays = tuple(a._array for a in arrays)%0A
     # Note: 
@@ -458,16 +458,29 @@
  return 
+ndarray._new(
 np.conca
@@ -500,24 +500,25 @@
 , axis=axis)
+)
 %0A%0Adef expand
@@ -713,16 +713,29 @@
  return 
+ndarray._new(
 np.expan
@@ -762,15 +762,23 @@
 on(x
+._array
 , axis)
+)
 %0A%0Ade
@@ -1004,16 +1004,29 @@
  return 
+ndarray._new(
 np.flip.
@@ -1034,32 +1034,39 @@
 implementation(x
+._array
 , axis=axis)%0A%0Ade
@@ -1057,24 +1057,25 @@
 , axis=axis)
+)
 %0A%0Adef reshap
@@ -1267,24 +1267,37 @@
 %0A    return 
+ndarray._new(
 np.reshape._
@@ -1316,16 +1316,23 @@
 on(x
+._array
 , shape)
 %0A%0Ade
@@ -1327,16 +1327,17 @@
 , shape)
+)
 %0A%0Adef ro
@@ -1595,16 +1595,29 @@
  return 
+ndarray._new(
 np.roll.
@@ -1633,16 +1633,23 @@
 tation(x
+._array
 , shift,
@@ -1655,24 +1655,25 @@
 , axis=axis)
+)
 %0A%0Adef squeez
@@ -1896,24 +1896,39 @@
 %0A    return 
+ndarray._array(
 np.squeeze._
@@ -1943,16 +1943,23 @@
 tation(x
+._array
 , axis=a
@@ -1962,16 +1962,17 @@
 is=axis)
+)
 %0A%0Adef st
@@ -2165,23 +2165,83 @@
 %22%22%22%0A    
-return 
+arrays = tuple(a._array for a in arrays)%0A    return ndarray._array(
 np.stack
@@ -2275,9 +2275,10 @@
 is=axis)
+)
 %0A

Make the array API manipulation functions use the array API ndarray object