diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__config__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__config__.py new file mode 100644 index 0000000000000000000000000000000000000000..361cf053ddf1cf04a73f117f6bcffc7f928f6349 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__config__.py @@ -0,0 +1,162 @@ +# This file is generated by numpy's build process +# It contains system_info results at the time of building this package. +from enum import Enum +from numpy.core._multiarray_umath import ( + __cpu_features__, + __cpu_baseline__, + __cpu_dispatch__, +) + +__all__ = ["show"] +_built_with_meson = True + + +class DisplayModes(Enum): + stdout = "stdout" + dicts = "dicts" + + +def _cleanup(d): + """ + Removes empty values in a `dict` recursively + This ensures we remove values that Meson could not provide to CONFIG + """ + if isinstance(d, dict): + return {k: _cleanup(v) for k, v in d.items() if v and _cleanup(v)} + else: + return d + + +CONFIG = _cleanup( + { + "Compilers": { + "c": { + "name": "gcc", + "linker": r"ld.bfd", + "version": "10.2.1", + "commands": r"cc", + "args": r"-fno-strict-aliasing", + "linker args": r"-Wl,--strip-debug, -fno-strict-aliasing", + }, + "cython": { + "name": "cython", + "linker": r"cython", + "version": "3.0.8", + "commands": r"cython", + "args": r"", + "linker args": r"", + }, + "c++": { + "name": "gcc", + "linker": r"ld.bfd", + "version": "10.2.1", + "commands": r"c++", + "args": r"", + "linker args": r"-Wl,--strip-debug", + }, + }, + "Machine Information": { + "host": { + "cpu": "x86_64", + "family": "x86_64", + "endian": "little", + "system": "linux", + }, + "build": { + "cpu": "x86_64", + "family": "x86_64", + "endian": "little", + "system": "linux", + }, + "cross-compiled": bool("False".lower().replace("false", "")), + }, + "Build Dependencies": { + "blas": { + "name": "openblas64", + "found": bool("True".lower().replace("false", "")), + "version": "0.3.23.dev", + "detection method": "pkgconfig", + "include directory": r"/usr/local/include", + "lib directory": r"/usr/local/lib", + "openblas configuration": r"USE_64BITINT=1 DYNAMIC_ARCH=1 DYNAMIC_OLDER= NO_CBLAS= NO_LAPACK= NO_LAPACKE= NO_AFFINITY=1 USE_OPENMP= HASWELL MAX_THREADS=2", + "pc file directory": r"/usr/local/lib/pkgconfig", + }, + "lapack": { + "name": "dep140551260102944", + "found": bool("True".lower().replace("false", "")), + "version": "1.26.4", + "detection method": "internal", + "include directory": r"unknown", + "lib directory": r"unknown", + "openblas configuration": r"unknown", + "pc file directory": r"unknown", + }, + }, + "Python Information": { + "path": r"/opt/python/cp312-cp312/bin/python", + "version": "3.12", + }, + "SIMD Extensions": { + "baseline": __cpu_baseline__, + "found": [ + feature for feature in __cpu_dispatch__ if __cpu_features__[feature] + ], + "not found": [ + feature for feature in __cpu_dispatch__ if not __cpu_features__[feature] + ], + }, + } +) + + +def _check_pyyaml(): + import yaml + + return yaml + + +def show(mode=DisplayModes.stdout.value): + """ + Show libraries and system information on which NumPy was built + and is being used + + Parameters + ---------- + mode : {`'stdout'`, `'dicts'`}, optional. + Indicates how to display the config information. + `'stdout'` prints to console, `'dicts'` returns a dictionary + of the configuration. + + Returns + ------- + out : {`dict`, `None`} + If mode is `'dicts'`, a dict is returned, else None + + See Also + -------- + get_include : Returns the directory containing NumPy C + header files. + + Notes + ----- + 1. The `'stdout'` mode will give more readable + output if ``pyyaml`` is installed + + """ + if mode == DisplayModes.stdout.value: + try: # Non-standard library, check import + yaml = _check_pyyaml() + + print(yaml.dump(CONFIG)) + except ModuleNotFoundError: + import warnings + import json + + warnings.warn("Install `pyyaml` for better output", stacklevel=1) + print(json.dumps(CONFIG, indent=2)) + elif mode == DisplayModes.dicts.value: + return CONFIG + else: + raise AttributeError( + f"Invalid `mode`, use one of: {', '.join([e.value for e in DisplayModes])}" + ) diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__init__.cython-30.pxd b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__init__.cython-30.pxd new file mode 100644 index 0000000000000000000000000000000000000000..1409514f7a845501a7787f6acb3a5570502d330d --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__init__.cython-30.pxd @@ -0,0 +1,1050 @@ +# NumPy static imports for Cython >= 3.0 +# +# If any of the PyArray_* functions are called, import_array must be +# called first. This is done automatically by Cython 3.0+ if a call +# is not detected inside of the module. +# +# Author: Dag Sverre Seljebotn +# + +from cpython.ref cimport Py_INCREF +from cpython.object cimport PyObject, PyTypeObject, PyObject_TypeCheck +cimport libc.stdio as stdio + + +cdef extern from *: + # Leave a marker that the NumPy declarations came from NumPy itself and not from Cython. + # See https://github.com/cython/cython/issues/3573 + """ + /* Using NumPy API declarations from "numpy/__init__.cython-30.pxd" */ + """ + + +cdef extern from "Python.h": + ctypedef int Py_intptr_t + +cdef extern from "numpy/arrayobject.h": + ctypedef Py_intptr_t npy_intp + ctypedef size_t npy_uintp + + cdef enum NPY_TYPES: + NPY_BOOL + NPY_BYTE + NPY_UBYTE + NPY_SHORT + NPY_USHORT + NPY_INT + NPY_UINT + NPY_LONG + NPY_ULONG + NPY_LONGLONG + NPY_ULONGLONG + NPY_FLOAT + NPY_DOUBLE + NPY_LONGDOUBLE + NPY_CFLOAT + NPY_CDOUBLE + NPY_CLONGDOUBLE + NPY_OBJECT + NPY_STRING + NPY_UNICODE + NPY_VOID + NPY_DATETIME + NPY_TIMEDELTA + NPY_NTYPES + NPY_NOTYPE + + NPY_INT8 + NPY_INT16 + NPY_INT32 + NPY_INT64 + NPY_INT128 + NPY_INT256 + NPY_UINT8 + NPY_UINT16 + NPY_UINT32 + NPY_UINT64 + NPY_UINT128 + NPY_UINT256 + NPY_FLOAT16 + NPY_FLOAT32 + NPY_FLOAT64 + NPY_FLOAT80 + NPY_FLOAT96 + NPY_FLOAT128 + NPY_FLOAT256 + NPY_COMPLEX32 + NPY_COMPLEX64 + NPY_COMPLEX128 + NPY_COMPLEX160 + NPY_COMPLEX192 + NPY_COMPLEX256 + NPY_COMPLEX512 + + NPY_INTP + + ctypedef enum NPY_ORDER: + NPY_ANYORDER + NPY_CORDER + NPY_FORTRANORDER + NPY_KEEPORDER + + ctypedef enum NPY_CASTING: + NPY_NO_CASTING + NPY_EQUIV_CASTING + NPY_SAFE_CASTING + NPY_SAME_KIND_CASTING + NPY_UNSAFE_CASTING + + ctypedef enum NPY_CLIPMODE: + NPY_CLIP + NPY_WRAP + NPY_RAISE + + ctypedef enum NPY_SCALARKIND: + NPY_NOSCALAR, + NPY_BOOL_SCALAR, + NPY_INTPOS_SCALAR, + NPY_INTNEG_SCALAR, + NPY_FLOAT_SCALAR, + NPY_COMPLEX_SCALAR, + NPY_OBJECT_SCALAR + + ctypedef enum NPY_SORTKIND: + NPY_QUICKSORT + NPY_HEAPSORT + NPY_MERGESORT + + ctypedef enum NPY_SEARCHSIDE: + NPY_SEARCHLEFT + NPY_SEARCHRIGHT + + enum: + # DEPRECATED since NumPy 1.7 ! Do not use in new code! + NPY_C_CONTIGUOUS + NPY_F_CONTIGUOUS + NPY_CONTIGUOUS + NPY_FORTRAN + NPY_OWNDATA + NPY_FORCECAST + NPY_ENSURECOPY + NPY_ENSUREARRAY + NPY_ELEMENTSTRIDES + NPY_ALIGNED + NPY_NOTSWAPPED + NPY_WRITEABLE + NPY_ARR_HAS_DESCR + + NPY_BEHAVED + NPY_BEHAVED_NS + NPY_CARRAY + NPY_CARRAY_RO + NPY_FARRAY + NPY_FARRAY_RO + NPY_DEFAULT + + NPY_IN_ARRAY + NPY_OUT_ARRAY + NPY_INOUT_ARRAY + NPY_IN_FARRAY + NPY_OUT_FARRAY + NPY_INOUT_FARRAY + + NPY_UPDATE_ALL + + enum: + # Added in NumPy 1.7 to replace the deprecated enums above. + NPY_ARRAY_C_CONTIGUOUS + NPY_ARRAY_F_CONTIGUOUS + NPY_ARRAY_OWNDATA + NPY_ARRAY_FORCECAST + NPY_ARRAY_ENSURECOPY + NPY_ARRAY_ENSUREARRAY + NPY_ARRAY_ELEMENTSTRIDES + NPY_ARRAY_ALIGNED + NPY_ARRAY_NOTSWAPPED + NPY_ARRAY_WRITEABLE + NPY_ARRAY_WRITEBACKIFCOPY + + NPY_ARRAY_BEHAVED + NPY_ARRAY_BEHAVED_NS + NPY_ARRAY_CARRAY + NPY_ARRAY_CARRAY_RO + NPY_ARRAY_FARRAY + NPY_ARRAY_FARRAY_RO + NPY_ARRAY_DEFAULT + + NPY_ARRAY_IN_ARRAY + NPY_ARRAY_OUT_ARRAY + NPY_ARRAY_INOUT_ARRAY + NPY_ARRAY_IN_FARRAY + NPY_ARRAY_OUT_FARRAY + NPY_ARRAY_INOUT_FARRAY + + NPY_ARRAY_UPDATE_ALL + + cdef enum: + NPY_MAXDIMS + + npy_intp NPY_MAX_ELSIZE + + ctypedef void (*PyArray_VectorUnaryFunc)(void *, void *, npy_intp, void *, void *) + + ctypedef struct PyArray_ArrayDescr: + # shape is a tuple, but Cython doesn't support "tuple shape" + # inside a non-PyObject declaration, so we have to declare it + # as just a PyObject*. + PyObject* shape + + ctypedef struct PyArray_Descr: + pass + + ctypedef class numpy.dtype [object PyArray_Descr, check_size ignore]: + # Use PyDataType_* macros when possible, however there are no macros + # for accessing some of the fields, so some are defined. + cdef PyTypeObject* typeobj + cdef char kind + cdef char type + # Numpy sometimes mutates this without warning (e.g. it'll + # sometimes change "|" to "<" in shared dtype objects on + # little-endian machines). If this matters to you, use + # PyArray_IsNativeByteOrder(dtype.byteorder) instead of + # directly accessing this field. + cdef char byteorder + cdef char flags + cdef int type_num + cdef int itemsize "elsize" + cdef int alignment + cdef object fields + cdef tuple names + # Use PyDataType_HASSUBARRAY to test whether this field is + # valid (the pointer can be NULL). Most users should access + # this field via the inline helper method PyDataType_SHAPE. + cdef PyArray_ArrayDescr* subarray + + ctypedef class numpy.flatiter [object PyArrayIterObject, check_size ignore]: + # Use through macros + pass + + ctypedef class numpy.broadcast [object PyArrayMultiIterObject, check_size ignore]: + # Use through macros + pass + + ctypedef struct PyArrayObject: + # For use in situations where ndarray can't replace PyArrayObject*, + # like PyArrayObject**. + pass + + ctypedef class numpy.ndarray [object PyArrayObject, check_size ignore]: + cdef __cythonbufferdefaults__ = {"mode": "strided"} + + # NOTE: no field declarations since direct access is deprecated since NumPy 1.7 + # Instead, we use properties that map to the corresponding C-API functions. + + @property + cdef inline PyObject* base(self) nogil: + """Returns a borrowed reference to the object owning the data/memory. + """ + return PyArray_BASE(self) + + @property + cdef inline dtype descr(self): + """Returns an owned reference to the dtype of the array. + """ + return PyArray_DESCR(self) + + @property + cdef inline int ndim(self) nogil: + """Returns the number of dimensions in the array. + """ + return PyArray_NDIM(self) + + @property + cdef inline npy_intp *shape(self) nogil: + """Returns a pointer to the dimensions/shape of the array. + The number of elements matches the number of dimensions of the array (ndim). + Can return NULL for 0-dimensional arrays. + """ + return PyArray_DIMS(self) + + @property + cdef inline npy_intp *strides(self) nogil: + """Returns a pointer to the strides of the array. + The number of elements matches the number of dimensions of the array (ndim). + """ + return PyArray_STRIDES(self) + + @property + cdef inline npy_intp size(self) nogil: + """Returns the total size (in number of elements) of the array. + """ + return PyArray_SIZE(self) + + @property + cdef inline char* data(self) nogil: + """The pointer to the data buffer as a char*. + This is provided for legacy reasons to avoid direct struct field access. + For new code that needs this access, you probably want to cast the result + of `PyArray_DATA()` instead, which returns a 'void*'. + """ + return PyArray_BYTES(self) + + ctypedef unsigned char npy_bool + + ctypedef signed char npy_byte + ctypedef signed short npy_short + ctypedef signed int npy_int + ctypedef signed long npy_long + ctypedef signed long long npy_longlong + + ctypedef unsigned char npy_ubyte + ctypedef unsigned short npy_ushort + ctypedef unsigned int npy_uint + ctypedef unsigned long npy_ulong + ctypedef unsigned long long npy_ulonglong + + ctypedef float npy_float + ctypedef double npy_double + ctypedef long double npy_longdouble + + ctypedef signed char npy_int8 + ctypedef signed short npy_int16 + ctypedef signed int npy_int32 + ctypedef signed long long npy_int64 + ctypedef signed long long npy_int96 + ctypedef signed long long npy_int128 + + ctypedef unsigned char npy_uint8 + ctypedef unsigned short npy_uint16 + ctypedef unsigned int npy_uint32 + ctypedef unsigned long long npy_uint64 + ctypedef unsigned long long npy_uint96 + ctypedef unsigned long long npy_uint128 + + ctypedef float npy_float32 + ctypedef double npy_float64 + ctypedef long double npy_float80 + ctypedef long double npy_float96 + ctypedef long double npy_float128 + + ctypedef struct npy_cfloat: + float real + float imag + + ctypedef struct npy_cdouble: + double real + double imag + + ctypedef struct npy_clongdouble: + long double real + long double imag + + ctypedef struct npy_complex64: + float real + float imag + + ctypedef struct npy_complex128: + double real + double imag + + ctypedef struct npy_complex160: + long double real + long double imag + + ctypedef struct npy_complex192: + long double real + long double imag + + ctypedef struct npy_complex256: + long double real + long double imag + + ctypedef struct PyArray_Dims: + npy_intp *ptr + int len + + int _import_array() except -1 + # A second definition so _import_array isn't marked as used when we use it here. + # Do not use - subject to change any time. + int __pyx_import_array "_import_array"() except -1 + + # + # Macros from ndarrayobject.h + # + bint PyArray_CHKFLAGS(ndarray m, int flags) nogil + bint PyArray_IS_C_CONTIGUOUS(ndarray arr) nogil + bint PyArray_IS_F_CONTIGUOUS(ndarray arr) nogil + bint PyArray_ISCONTIGUOUS(ndarray m) nogil + bint PyArray_ISWRITEABLE(ndarray m) nogil + bint PyArray_ISALIGNED(ndarray m) nogil + + int PyArray_NDIM(ndarray) nogil + bint PyArray_ISONESEGMENT(ndarray) nogil + bint PyArray_ISFORTRAN(ndarray) nogil + int PyArray_FORTRANIF(ndarray) nogil + + void* PyArray_DATA(ndarray) nogil + char* PyArray_BYTES(ndarray) nogil + + npy_intp* PyArray_DIMS(ndarray) nogil + npy_intp* PyArray_STRIDES(ndarray) nogil + npy_intp PyArray_DIM(ndarray, size_t) nogil + npy_intp PyArray_STRIDE(ndarray, size_t) nogil + + PyObject *PyArray_BASE(ndarray) nogil # returns borrowed reference! + PyArray_Descr *PyArray_DESCR(ndarray) nogil # returns borrowed reference to dtype! + PyArray_Descr *PyArray_DTYPE(ndarray) nogil # returns borrowed reference to dtype! NP 1.7+ alias for descr. + int PyArray_FLAGS(ndarray) nogil + void PyArray_CLEARFLAGS(ndarray, int flags) nogil # Added in NumPy 1.7 + void PyArray_ENABLEFLAGS(ndarray, int flags) nogil # Added in NumPy 1.7 + npy_intp PyArray_ITEMSIZE(ndarray) nogil + int PyArray_TYPE(ndarray arr) nogil + + object PyArray_GETITEM(ndarray arr, void *itemptr) + int PyArray_SETITEM(ndarray arr, void *itemptr, object obj) except -1 + + bint PyTypeNum_ISBOOL(int) nogil + bint PyTypeNum_ISUNSIGNED(int) nogil + bint PyTypeNum_ISSIGNED(int) nogil + bint PyTypeNum_ISINTEGER(int) nogil + bint PyTypeNum_ISFLOAT(int) nogil + bint PyTypeNum_ISNUMBER(int) nogil + bint PyTypeNum_ISSTRING(int) nogil + bint PyTypeNum_ISCOMPLEX(int) nogil + bint PyTypeNum_ISPYTHON(int) nogil + bint PyTypeNum_ISFLEXIBLE(int) nogil + bint PyTypeNum_ISUSERDEF(int) nogil + bint PyTypeNum_ISEXTENDED(int) nogil + bint PyTypeNum_ISOBJECT(int) nogil + + bint PyDataType_ISBOOL(dtype) nogil + bint PyDataType_ISUNSIGNED(dtype) nogil + bint PyDataType_ISSIGNED(dtype) nogil + bint PyDataType_ISINTEGER(dtype) nogil + bint PyDataType_ISFLOAT(dtype) nogil + bint PyDataType_ISNUMBER(dtype) nogil + bint PyDataType_ISSTRING(dtype) nogil + bint PyDataType_ISCOMPLEX(dtype) nogil + bint PyDataType_ISPYTHON(dtype) nogil + bint PyDataType_ISFLEXIBLE(dtype) nogil + bint PyDataType_ISUSERDEF(dtype) nogil + bint PyDataType_ISEXTENDED(dtype) nogil + bint PyDataType_ISOBJECT(dtype) nogil + bint PyDataType_HASFIELDS(dtype) nogil + bint PyDataType_HASSUBARRAY(dtype) nogil + + bint PyArray_ISBOOL(ndarray) nogil + bint PyArray_ISUNSIGNED(ndarray) nogil + bint PyArray_ISSIGNED(ndarray) nogil + bint PyArray_ISINTEGER(ndarray) nogil + bint PyArray_ISFLOAT(ndarray) nogil + bint PyArray_ISNUMBER(ndarray) nogil + bint PyArray_ISSTRING(ndarray) nogil + bint PyArray_ISCOMPLEX(ndarray) nogil + bint PyArray_ISPYTHON(ndarray) nogil + bint PyArray_ISFLEXIBLE(ndarray) nogil + bint PyArray_ISUSERDEF(ndarray) nogil + bint PyArray_ISEXTENDED(ndarray) nogil + bint PyArray_ISOBJECT(ndarray) nogil + bint PyArray_HASFIELDS(ndarray) nogil + + bint PyArray_ISVARIABLE(ndarray) nogil + + bint PyArray_SAFEALIGNEDCOPY(ndarray) nogil + bint PyArray_ISNBO(char) nogil # works on ndarray.byteorder + bint PyArray_IsNativeByteOrder(char) nogil # works on ndarray.byteorder + bint PyArray_ISNOTSWAPPED(ndarray) nogil + bint PyArray_ISBYTESWAPPED(ndarray) nogil + + bint PyArray_FLAGSWAP(ndarray, int) nogil + + bint PyArray_ISCARRAY(ndarray) nogil + bint PyArray_ISCARRAY_RO(ndarray) nogil + bint PyArray_ISFARRAY(ndarray) nogil + bint PyArray_ISFARRAY_RO(ndarray) nogil + bint PyArray_ISBEHAVED(ndarray) nogil + bint PyArray_ISBEHAVED_RO(ndarray) nogil + + + bint PyDataType_ISNOTSWAPPED(dtype) nogil + bint PyDataType_ISBYTESWAPPED(dtype) nogil + + bint PyArray_DescrCheck(object) + + bint PyArray_Check(object) + bint PyArray_CheckExact(object) + + # Cannot be supported due to out arg: + # bint PyArray_HasArrayInterfaceType(object, dtype, object, object&) + # bint PyArray_HasArrayInterface(op, out) + + + bint PyArray_IsZeroDim(object) + # Cannot be supported due to ## ## in macro: + # bint PyArray_IsScalar(object, verbatim work) + bint PyArray_CheckScalar(object) + bint PyArray_IsPythonNumber(object) + bint PyArray_IsPythonScalar(object) + bint PyArray_IsAnyScalar(object) + bint PyArray_CheckAnyScalar(object) + + ndarray PyArray_GETCONTIGUOUS(ndarray) + bint PyArray_SAMESHAPE(ndarray, ndarray) nogil + npy_intp PyArray_SIZE(ndarray) nogil + npy_intp PyArray_NBYTES(ndarray) nogil + + object PyArray_FROM_O(object) + object PyArray_FROM_OF(object m, int flags) + object PyArray_FROM_OT(object m, int type) + object PyArray_FROM_OTF(object m, int type, int flags) + object PyArray_FROMANY(object m, int type, int min, int max, int flags) + object PyArray_ZEROS(int nd, npy_intp* dims, int type, int fortran) + object PyArray_EMPTY(int nd, npy_intp* dims, int type, int fortran) + void PyArray_FILLWBYTE(object, int val) + npy_intp PyArray_REFCOUNT(object) + object PyArray_ContiguousFromAny(op, int, int min_depth, int max_depth) + unsigned char PyArray_EquivArrTypes(ndarray a1, ndarray a2) + bint PyArray_EquivByteorders(int b1, int b2) nogil + object PyArray_SimpleNew(int nd, npy_intp* dims, int typenum) + object PyArray_SimpleNewFromData(int nd, npy_intp* dims, int typenum, void* data) + #object PyArray_SimpleNewFromDescr(int nd, npy_intp* dims, dtype descr) + object PyArray_ToScalar(void* data, ndarray arr) + + void* PyArray_GETPTR1(ndarray m, npy_intp i) nogil + void* PyArray_GETPTR2(ndarray m, npy_intp i, npy_intp j) nogil + void* PyArray_GETPTR3(ndarray m, npy_intp i, npy_intp j, npy_intp k) nogil + void* PyArray_GETPTR4(ndarray m, npy_intp i, npy_intp j, npy_intp k, npy_intp l) nogil + + # Cannot be supported due to out arg + # void PyArray_DESCR_REPLACE(descr) + + + object PyArray_Copy(ndarray) + object PyArray_FromObject(object op, int type, int min_depth, int max_depth) + object PyArray_ContiguousFromObject(object op, int type, int min_depth, int max_depth) + object PyArray_CopyFromObject(object op, int type, int min_depth, int max_depth) + + object PyArray_Cast(ndarray mp, int type_num) + object PyArray_Take(ndarray ap, object items, int axis) + object PyArray_Put(ndarray ap, object items, object values) + + void PyArray_ITER_RESET(flatiter it) nogil + void PyArray_ITER_NEXT(flatiter it) nogil + void PyArray_ITER_GOTO(flatiter it, npy_intp* destination) nogil + void PyArray_ITER_GOTO1D(flatiter it, npy_intp ind) nogil + void* PyArray_ITER_DATA(flatiter it) nogil + bint PyArray_ITER_NOTDONE(flatiter it) nogil + + void PyArray_MultiIter_RESET(broadcast multi) nogil + void PyArray_MultiIter_NEXT(broadcast multi) nogil + void PyArray_MultiIter_GOTO(broadcast multi, npy_intp dest) nogil + void PyArray_MultiIter_GOTO1D(broadcast multi, npy_intp ind) nogil + void* PyArray_MultiIter_DATA(broadcast multi, npy_intp i) nogil + void PyArray_MultiIter_NEXTi(broadcast multi, npy_intp i) nogil + bint PyArray_MultiIter_NOTDONE(broadcast multi) nogil + + # Functions from __multiarray_api.h + + # Functions taking dtype and returning object/ndarray are disabled + # for now as they steal dtype references. I'm conservative and disable + # more than is probably needed until it can be checked further. + int PyArray_SetNumericOps (object) except -1 + object PyArray_GetNumericOps () + int PyArray_INCREF (ndarray) except * # uses PyArray_Item_INCREF... + int PyArray_XDECREF (ndarray) except * # uses PyArray_Item_DECREF... + void PyArray_SetStringFunction (object, int) + dtype PyArray_DescrFromType (int) + object PyArray_TypeObjectFromType (int) + char * PyArray_Zero (ndarray) + char * PyArray_One (ndarray) + #object PyArray_CastToType (ndarray, dtype, int) + int PyArray_CastTo (ndarray, ndarray) except -1 + int PyArray_CastAnyTo (ndarray, ndarray) except -1 + int PyArray_CanCastSafely (int, int) # writes errors + npy_bool PyArray_CanCastTo (dtype, dtype) # writes errors + int PyArray_ObjectType (object, int) except 0 + dtype PyArray_DescrFromObject (object, dtype) + #ndarray* PyArray_ConvertToCommonType (object, int *) + dtype PyArray_DescrFromScalar (object) + dtype PyArray_DescrFromTypeObject (object) + npy_intp PyArray_Size (object) + #object PyArray_Scalar (void *, dtype, object) + #object PyArray_FromScalar (object, dtype) + void PyArray_ScalarAsCtype (object, void *) + #int PyArray_CastScalarToCtype (object, void *, dtype) + #int PyArray_CastScalarDirect (object, dtype, void *, int) + object PyArray_ScalarFromObject (object) + #PyArray_VectorUnaryFunc * PyArray_GetCastFunc (dtype, int) + object PyArray_FromDims (int, int *, int) + #object PyArray_FromDimsAndDataAndDescr (int, int *, dtype, char *) + #object PyArray_FromAny (object, dtype, int, int, int, object) + object PyArray_EnsureArray (object) + object PyArray_EnsureAnyArray (object) + #object PyArray_FromFile (stdio.FILE *, dtype, npy_intp, char *) + #object PyArray_FromString (char *, npy_intp, dtype, npy_intp, char *) + #object PyArray_FromBuffer (object, dtype, npy_intp, npy_intp) + #object PyArray_FromIter (object, dtype, npy_intp) + object PyArray_Return (ndarray) + #object PyArray_GetField (ndarray, dtype, int) + #int PyArray_SetField (ndarray, dtype, int, object) except -1 + object PyArray_Byteswap (ndarray, npy_bool) + object PyArray_Resize (ndarray, PyArray_Dims *, int, NPY_ORDER) + int PyArray_MoveInto (ndarray, ndarray) except -1 + int PyArray_CopyInto (ndarray, ndarray) except -1 + int PyArray_CopyAnyInto (ndarray, ndarray) except -1 + int PyArray_CopyObject (ndarray, object) except -1 + object PyArray_NewCopy (ndarray, NPY_ORDER) + object PyArray_ToList (ndarray) + object PyArray_ToString (ndarray, NPY_ORDER) + int PyArray_ToFile (ndarray, stdio.FILE *, char *, char *) except -1 + int PyArray_Dump (object, object, int) except -1 + object PyArray_Dumps (object, int) + int PyArray_ValidType (int) # Cannot error + void PyArray_UpdateFlags (ndarray, int) + object PyArray_New (type, int, npy_intp *, int, npy_intp *, void *, int, int, object) + #object PyArray_NewFromDescr (type, dtype, int, npy_intp *, npy_intp *, void *, int, object) + #dtype PyArray_DescrNew (dtype) + dtype PyArray_DescrNewFromType (int) + double PyArray_GetPriority (object, double) # clears errors as of 1.25 + object PyArray_IterNew (object) + object PyArray_MultiIterNew (int, ...) + + int PyArray_PyIntAsInt (object) except? -1 + npy_intp PyArray_PyIntAsIntp (object) + int PyArray_Broadcast (broadcast) except -1 + void PyArray_FillObjectArray (ndarray, object) except * + int PyArray_FillWithScalar (ndarray, object) except -1 + npy_bool PyArray_CheckStrides (int, int, npy_intp, npy_intp, npy_intp *, npy_intp *) + dtype PyArray_DescrNewByteorder (dtype, char) + object PyArray_IterAllButAxis (object, int *) + #object PyArray_CheckFromAny (object, dtype, int, int, int, object) + #object PyArray_FromArray (ndarray, dtype, int) + object PyArray_FromInterface (object) + object PyArray_FromStructInterface (object) + #object PyArray_FromArrayAttr (object, dtype, object) + #NPY_SCALARKIND PyArray_ScalarKind (int, ndarray*) + int PyArray_CanCoerceScalar (int, int, NPY_SCALARKIND) + object PyArray_NewFlagsObject (object) + npy_bool PyArray_CanCastScalar (type, type) + #int PyArray_CompareUCS4 (npy_ucs4 *, npy_ucs4 *, register size_t) + int PyArray_RemoveSmallest (broadcast) except -1 + int PyArray_ElementStrides (object) + void PyArray_Item_INCREF (char *, dtype) except * + void PyArray_Item_XDECREF (char *, dtype) except * + object PyArray_FieldNames (object) + object PyArray_Transpose (ndarray, PyArray_Dims *) + object PyArray_TakeFrom (ndarray, object, int, ndarray, NPY_CLIPMODE) + object PyArray_PutTo (ndarray, object, object, NPY_CLIPMODE) + object PyArray_PutMask (ndarray, object, object) + object PyArray_Repeat (ndarray, object, int) + object PyArray_Choose (ndarray, object, ndarray, NPY_CLIPMODE) + int PyArray_Sort (ndarray, int, NPY_SORTKIND) except -1 + object PyArray_ArgSort (ndarray, int, NPY_SORTKIND) + object PyArray_SearchSorted (ndarray, object, NPY_SEARCHSIDE, PyObject *) + object PyArray_ArgMax (ndarray, int, ndarray) + object PyArray_ArgMin (ndarray, int, ndarray) + object PyArray_Reshape (ndarray, object) + object PyArray_Newshape (ndarray, PyArray_Dims *, NPY_ORDER) + object PyArray_Squeeze (ndarray) + #object PyArray_View (ndarray, dtype, type) + object PyArray_SwapAxes (ndarray, int, int) + object PyArray_Max (ndarray, int, ndarray) + object PyArray_Min (ndarray, int, ndarray) + object PyArray_Ptp (ndarray, int, ndarray) + object PyArray_Mean (ndarray, int, int, ndarray) + object PyArray_Trace (ndarray, int, int, int, int, ndarray) + object PyArray_Diagonal (ndarray, int, int, int) + object PyArray_Clip (ndarray, object, object, ndarray) + object PyArray_Conjugate (ndarray, ndarray) + object PyArray_Nonzero (ndarray) + object PyArray_Std (ndarray, int, int, ndarray, int) + object PyArray_Sum (ndarray, int, int, ndarray) + object PyArray_CumSum (ndarray, int, int, ndarray) + object PyArray_Prod (ndarray, int, int, ndarray) + object PyArray_CumProd (ndarray, int, int, ndarray) + object PyArray_All (ndarray, int, ndarray) + object PyArray_Any (ndarray, int, ndarray) + object PyArray_Compress (ndarray, object, int, ndarray) + object PyArray_Flatten (ndarray, NPY_ORDER) + object PyArray_Ravel (ndarray, NPY_ORDER) + npy_intp PyArray_MultiplyList (npy_intp *, int) + int PyArray_MultiplyIntList (int *, int) + void * PyArray_GetPtr (ndarray, npy_intp*) + int PyArray_CompareLists (npy_intp *, npy_intp *, int) + #int PyArray_AsCArray (object*, void *, npy_intp *, int, dtype) + #int PyArray_As1D (object*, char **, int *, int) + #int PyArray_As2D (object*, char ***, int *, int *, int) + int PyArray_Free (object, void *) + #int PyArray_Converter (object, object*) + int PyArray_IntpFromSequence (object, npy_intp *, int) except -1 + object PyArray_Concatenate (object, int) + object PyArray_InnerProduct (object, object) + object PyArray_MatrixProduct (object, object) + object PyArray_CopyAndTranspose (object) + object PyArray_Correlate (object, object, int) + int PyArray_TypestrConvert (int, int) + #int PyArray_DescrConverter (object, dtype*) except 0 + #int PyArray_DescrConverter2 (object, dtype*) except 0 + int PyArray_IntpConverter (object, PyArray_Dims *) except 0 + #int PyArray_BufferConverter (object, chunk) except 0 + int PyArray_AxisConverter (object, int *) except 0 + int PyArray_BoolConverter (object, npy_bool *) except 0 + int PyArray_ByteorderConverter (object, char *) except 0 + int PyArray_OrderConverter (object, NPY_ORDER *) except 0 + unsigned char PyArray_EquivTypes (dtype, dtype) # clears errors + #object PyArray_Zeros (int, npy_intp *, dtype, int) + #object PyArray_Empty (int, npy_intp *, dtype, int) + object PyArray_Where (object, object, object) + object PyArray_Arange (double, double, double, int) + #object PyArray_ArangeObj (object, object, object, dtype) + int PyArray_SortkindConverter (object, NPY_SORTKIND *) except 0 + object PyArray_LexSort (object, int) + object PyArray_Round (ndarray, int, ndarray) + unsigned char PyArray_EquivTypenums (int, int) + int PyArray_RegisterDataType (dtype) except -1 + int PyArray_RegisterCastFunc (dtype, int, PyArray_VectorUnaryFunc *) except -1 + int PyArray_RegisterCanCast (dtype, int, NPY_SCALARKIND) except -1 + #void PyArray_InitArrFuncs (PyArray_ArrFuncs *) + object PyArray_IntTupleFromIntp (int, npy_intp *) + int PyArray_TypeNumFromName (char *) + int PyArray_ClipmodeConverter (object, NPY_CLIPMODE *) except 0 + #int PyArray_OutputConverter (object, ndarray*) except 0 + object PyArray_BroadcastToShape (object, npy_intp *, int) + void _PyArray_SigintHandler (int) + void* _PyArray_GetSigintBuf () + #int PyArray_DescrAlignConverter (object, dtype*) except 0 + #int PyArray_DescrAlignConverter2 (object, dtype*) except 0 + int PyArray_SearchsideConverter (object, void *) except 0 + object PyArray_CheckAxis (ndarray, int *, int) + npy_intp PyArray_OverflowMultiplyList (npy_intp *, int) + int PyArray_CompareString (char *, char *, size_t) + int PyArray_SetBaseObject(ndarray, base) except -1 # NOTE: steals a reference to base! Use "set_array_base()" instead. + + +# Typedefs that matches the runtime dtype objects in +# the numpy module. + +# The ones that are commented out needs an IFDEF function +# in Cython to enable them only on the right systems. + +ctypedef npy_int8 int8_t +ctypedef npy_int16 int16_t +ctypedef npy_int32 int32_t +ctypedef npy_int64 int64_t +#ctypedef npy_int96 int96_t +#ctypedef npy_int128 int128_t + +ctypedef npy_uint8 uint8_t +ctypedef npy_uint16 uint16_t +ctypedef npy_uint32 uint32_t +ctypedef npy_uint64 uint64_t +#ctypedef npy_uint96 uint96_t +#ctypedef npy_uint128 uint128_t + +ctypedef npy_float32 float32_t +ctypedef npy_float64 float64_t +#ctypedef npy_float80 float80_t +#ctypedef npy_float128 float128_t + +ctypedef float complex complex64_t +ctypedef double complex complex128_t + +# The int types are mapped a bit surprising -- +# numpy.int corresponds to 'l' and numpy.long to 'q' +ctypedef npy_long int_t +ctypedef npy_longlong longlong_t + +ctypedef npy_ulong uint_t +ctypedef npy_ulonglong ulonglong_t + +ctypedef npy_intp intp_t +ctypedef npy_uintp uintp_t + +ctypedef npy_double float_t +ctypedef npy_double double_t +ctypedef npy_longdouble longdouble_t + +ctypedef npy_cfloat cfloat_t +ctypedef npy_cdouble cdouble_t +ctypedef npy_clongdouble clongdouble_t + +ctypedef npy_cdouble complex_t + +cdef inline object PyArray_MultiIterNew1(a): + return PyArray_MultiIterNew(1, a) + +cdef inline object PyArray_MultiIterNew2(a, b): + return PyArray_MultiIterNew(2, a, b) + +cdef inline object PyArray_MultiIterNew3(a, b, c): + return PyArray_MultiIterNew(3, a, b, c) + +cdef inline object PyArray_MultiIterNew4(a, b, c, d): + return PyArray_MultiIterNew(4, a, b, c, d) + +cdef inline object PyArray_MultiIterNew5(a, b, c, d, e): + return PyArray_MultiIterNew(5, a, b, c, d, e) + +cdef inline tuple PyDataType_SHAPE(dtype d): + if PyDataType_HASSUBARRAY(d): + return d.subarray.shape + else: + return () + + +cdef extern from "numpy/ndarrayobject.h": + PyTypeObject PyTimedeltaArrType_Type + PyTypeObject PyDatetimeArrType_Type + ctypedef int64_t npy_timedelta + ctypedef int64_t npy_datetime + +cdef extern from "numpy/ndarraytypes.h": + ctypedef struct PyArray_DatetimeMetaData: + NPY_DATETIMEUNIT base + int64_t num + +cdef extern from "numpy/arrayscalars.h": + + # abstract types + ctypedef class numpy.generic [object PyObject]: + pass + ctypedef class numpy.number [object PyObject]: + pass + ctypedef class numpy.integer [object PyObject]: + pass + ctypedef class numpy.signedinteger [object PyObject]: + pass + ctypedef class numpy.unsignedinteger [object PyObject]: + pass + ctypedef class numpy.inexact [object PyObject]: + pass + ctypedef class numpy.floating [object PyObject]: + pass + ctypedef class numpy.complexfloating [object PyObject]: + pass + ctypedef class numpy.flexible [object PyObject]: + pass + ctypedef class numpy.character [object PyObject]: + pass + + ctypedef struct PyDatetimeScalarObject: + # PyObject_HEAD + npy_datetime obval + PyArray_DatetimeMetaData obmeta + + ctypedef struct PyTimedeltaScalarObject: + # PyObject_HEAD + npy_timedelta obval + PyArray_DatetimeMetaData obmeta + + ctypedef enum NPY_DATETIMEUNIT: + NPY_FR_Y + NPY_FR_M + NPY_FR_W + NPY_FR_D + NPY_FR_B + NPY_FR_h + NPY_FR_m + NPY_FR_s + NPY_FR_ms + NPY_FR_us + NPY_FR_ns + NPY_FR_ps + NPY_FR_fs + NPY_FR_as + NPY_FR_GENERIC + + +# +# ufunc API +# + +cdef extern from "numpy/ufuncobject.h": + + ctypedef void (*PyUFuncGenericFunction) (char **, npy_intp *, npy_intp *, void *) + + ctypedef class numpy.ufunc [object PyUFuncObject, check_size ignore]: + cdef: + int nin, nout, nargs + int identity + PyUFuncGenericFunction *functions + void **data + int ntypes + int check_return + char *name + char *types + char *doc + void *ptr + PyObject *obj + PyObject *userloops + + cdef enum: + PyUFunc_Zero + PyUFunc_One + PyUFunc_None + UFUNC_ERR_IGNORE + UFUNC_ERR_WARN + UFUNC_ERR_RAISE + UFUNC_ERR_CALL + UFUNC_ERR_PRINT + UFUNC_ERR_LOG + UFUNC_MASK_DIVIDEBYZERO + UFUNC_MASK_OVERFLOW + UFUNC_MASK_UNDERFLOW + UFUNC_MASK_INVALID + UFUNC_SHIFT_DIVIDEBYZERO + UFUNC_SHIFT_OVERFLOW + UFUNC_SHIFT_UNDERFLOW + UFUNC_SHIFT_INVALID + UFUNC_FPE_DIVIDEBYZERO + UFUNC_FPE_OVERFLOW + UFUNC_FPE_UNDERFLOW + UFUNC_FPE_INVALID + UFUNC_ERR_DEFAULT + UFUNC_ERR_DEFAULT2 + + object PyUFunc_FromFuncAndData(PyUFuncGenericFunction *, + void **, char *, int, int, int, int, char *, char *, int) + int PyUFunc_RegisterLoopForType(ufunc, int, + PyUFuncGenericFunction, int *, void *) except -1 + void PyUFunc_f_f_As_d_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_d_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_f_f \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_g_g \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_F_F_As_D_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_F_F \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_D_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_G_G \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_O_O \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_ff_f_As_dd_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_ff_f \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_dd_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_gg_g \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_FF_F_As_DD_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_DD_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_FF_F \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_GG_G \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_OO_O \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_O_O_method \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_OO_O_method \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_On_Om \ + (char **, npy_intp *, npy_intp *, void *) + int PyUFunc_GetPyValues \ + (char *, int *, int *, PyObject **) + int PyUFunc_checkfperr \ + (int, PyObject *, int *) + void PyUFunc_clearfperr() + int PyUFunc_getfperr() + int PyUFunc_handlefperr \ + (int, PyObject *, int, int *) except -1 + int PyUFunc_ReplaceLoopBySignature \ + (ufunc, PyUFuncGenericFunction, int *, PyUFuncGenericFunction *) + object PyUFunc_FromFuncAndDataAndSignature \ + (PyUFuncGenericFunction *, void **, char *, int, int, int, + int, char *, char *, int, char *) + + int _import_umath() except -1 + +cdef inline void set_array_base(ndarray arr, object base): + Py_INCREF(base) # important to do this before stealing the reference below! + PyArray_SetBaseObject(arr, base) + +cdef inline object get_array_base(ndarray arr): + base = PyArray_BASE(arr) + if base is NULL: + return None + return base + +# Versions of the import_* functions which are more suitable for +# Cython code. +cdef inline int import_array() except -1: + try: + __pyx_import_array() + except Exception: + raise ImportError("numpy.core.multiarray failed to import") + +cdef inline int import_umath() except -1: + try: + _import_umath() + except Exception: + raise ImportError("numpy.core.umath failed to import") + +cdef inline int import_ufunc() except -1: + try: + _import_umath() + except Exception: + raise ImportError("numpy.core.umath failed to import") + + +cdef inline bint is_timedelta64_object(object obj): + """ + Cython equivalent of `isinstance(obj, np.timedelta64)` + + Parameters + ---------- + obj : object + + Returns + ------- + bool + """ + return PyObject_TypeCheck(obj, &PyTimedeltaArrType_Type) + + +cdef inline bint is_datetime64_object(object obj): + """ + Cython equivalent of `isinstance(obj, np.datetime64)` + + Parameters + ---------- + obj : object + + Returns + ------- + bool + """ + return PyObject_TypeCheck(obj, &PyDatetimeArrType_Type) + + +cdef inline npy_datetime get_datetime64_value(object obj) nogil: + """ + returns the int64 value underlying scalar numpy datetime64 object + + Note that to interpret this as a datetime, the corresponding unit is + also needed. That can be found using `get_datetime64_unit`. + """ + return (obj).obval + + +cdef inline npy_timedelta get_timedelta64_value(object obj) nogil: + """ + returns the int64 value underlying scalar numpy timedelta64 object + """ + return (obj).obval + + +cdef inline NPY_DATETIMEUNIT get_datetime64_unit(object obj) nogil: + """ + returns the unit part of the dtype for a numpy datetime64 object. + """ + return (obj).obmeta.base diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..91da496a95271f8e3eb4f4ea2cbaec925325a1f5 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__init__.py @@ -0,0 +1,461 @@ +""" +NumPy +===== + +Provides + 1. An array object of arbitrary homogeneous items + 2. Fast mathematical operations over arrays + 3. Linear Algebra, Fourier Transforms, Random Number Generation + +How to use the documentation +---------------------------- +Documentation is available in two forms: docstrings provided +with the code, and a loose standing reference guide, available from +`the NumPy homepage `_. + +We recommend exploring the docstrings using +`IPython `_, an advanced Python shell with +TAB-completion and introspection capabilities. See below for further +instructions. + +The docstring examples assume that `numpy` has been imported as ``np``:: + + >>> import numpy as np + +Code snippets are indicated by three greater-than signs:: + + >>> x = 42 + >>> x = x + 1 + +Use the built-in ``help`` function to view a function's docstring:: + + >>> help(np.sort) + ... # doctest: +SKIP + +For some objects, ``np.info(obj)`` may provide additional help. This is +particularly true if you see the line "Help on ufunc object:" at the top +of the help() page. Ufuncs are implemented in C, not Python, for speed. +The native Python help() does not know how to view their help, but our +np.info() function does. + +To search for documents containing a keyword, do:: + + >>> np.lookfor('keyword') + ... # doctest: +SKIP + +General-purpose documents like a glossary and help on the basic concepts +of numpy are available under the ``doc`` sub-module:: + + >>> from numpy import doc + >>> help(doc) + ... # doctest: +SKIP + +Available subpackages +--------------------- +lib + Basic functions used by several sub-packages. +random + Core Random Tools +linalg + Core Linear Algebra Tools +fft + Core FFT routines +polynomial + Polynomial tools +testing + NumPy testing tools +distutils + Enhancements to distutils with support for + Fortran compilers support and more (for Python <= 3.11). + +Utilities +--------- +test + Run numpy unittests +show_config + Show numpy build configuration +matlib + Make everything matrices. +__version__ + NumPy version string + +Viewing documentation using IPython +----------------------------------- + +Start IPython and import `numpy` usually under the alias ``np``: `import +numpy as np`. Then, directly past or use the ``%cpaste`` magic to paste +examples into the shell. To see which functions are available in `numpy`, +type ``np.`` (where ```` refers to the TAB key), or use +``np.*cos*?`` (where ```` refers to the ENTER key) to narrow +down the list. To view the docstring for a function, use +``np.cos?`` (to view the docstring) and ``np.cos??`` (to view +the source code). + +Copies vs. in-place operation +----------------------------- +Most of the functions in `numpy` return a copy of the array argument +(e.g., `np.sort`). In-place versions of these functions are often +available as array methods, i.e. ``x = np.array([1,2,3]); x.sort()``. +Exceptions to this rule are documented. + +""" +import sys +import warnings + +from ._globals import _NoValue, _CopyMode +# These exceptions were moved in 1.25 and are hidden from __dir__() +from .exceptions import ( + ComplexWarning, ModuleDeprecationWarning, VisibleDeprecationWarning, + TooHardError, AxisError) + + +# If a version with git hash was stored, use that instead +from . import version +from .version import __version__ + +# We first need to detect if we're being called as part of the numpy setup +# procedure itself in a reliable manner. +try: + __NUMPY_SETUP__ +except NameError: + __NUMPY_SETUP__ = False + +if __NUMPY_SETUP__: + sys.stderr.write('Running from numpy source directory.\n') +else: + # Allow distributors to run custom init code before importing numpy.core + from . import _distributor_init + + try: + from numpy.__config__ import show as show_config + except ImportError as e: + msg = """Error importing numpy: you should not try to import numpy from + its source directory; please exit the numpy source tree, and relaunch + your python interpreter from there.""" + raise ImportError(msg) from e + + __all__ = [ + 'exceptions', 'ModuleDeprecationWarning', 'VisibleDeprecationWarning', + 'ComplexWarning', 'TooHardError', 'AxisError'] + + # mapping of {name: (value, deprecation_msg)} + __deprecated_attrs__ = {} + + from . import core + from .core import * + from . import compat + from . import exceptions + from . import dtypes + from . import lib + # NOTE: to be revisited following future namespace cleanup. + # See gh-14454 and gh-15672 for discussion. + from .lib import * + + from . import linalg + from . import fft + from . import polynomial + from . import random + from . import ctypeslib + from . import ma + from . import matrixlib as _mat + from .matrixlib import * + + # Deprecations introduced in NumPy 1.20.0, 2020-06-06 + import builtins as _builtins + + _msg = ( + "module 'numpy' has no attribute '{n}'.\n" + "`np.{n}` was a deprecated alias for the builtin `{n}`. " + "To avoid this error in existing code, use `{n}` by itself. " + "Doing this will not modify any behavior and is safe. {extended_msg}\n" + "The aliases was originally deprecated in NumPy 1.20; for more " + "details and guidance see the original release note at:\n" + " https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations") + + _specific_msg = ( + "If you specifically wanted the numpy scalar type, use `np.{}` here.") + + _int_extended_msg = ( + "When replacing `np.{}`, you may wish to use e.g. `np.int64` " + "or `np.int32` to specify the precision. If you wish to review " + "your current use, check the release note link for " + "additional information.") + + _type_info = [ + ("object", ""), # The NumPy scalar only exists by name. + ("bool", _specific_msg.format("bool_")), + ("float", _specific_msg.format("float64")), + ("complex", _specific_msg.format("complex128")), + ("str", _specific_msg.format("str_")), + ("int", _int_extended_msg.format("int"))] + + __former_attrs__ = { + n: _msg.format(n=n, extended_msg=extended_msg) + for n, extended_msg in _type_info + } + + # Future warning introduced in NumPy 1.24.0, 2022-11-17 + _msg = ( + "`np.{n}` is a deprecated alias for `{an}`. (Deprecated NumPy 1.24)") + + # Some of these are awkward (since `np.str` may be preferable in the long + # term), but overall the names ending in 0 seem undesirable + _type_info = [ + ("bool8", bool_, "np.bool_"), + ("int0", intp, "np.intp"), + ("uint0", uintp, "np.uintp"), + ("str0", str_, "np.str_"), + ("bytes0", bytes_, "np.bytes_"), + ("void0", void, "np.void"), + ("object0", object_, + "`np.object0` is a deprecated alias for `np.object_`. " + "`object` can be used instead. (Deprecated NumPy 1.24)")] + + # Some of these could be defined right away, but most were aliases to + # the Python objects and only removed in NumPy 1.24. Defining them should + # probably wait for NumPy 1.26 or 2.0. + # When defined, these should possibly not be added to `__all__` to avoid + # import with `from numpy import *`. + __future_scalars__ = {"bool", "long", "ulong", "str", "bytes", "object"} + + __deprecated_attrs__.update({ + n: (alias, _msg.format(n=n, an=an)) for n, alias, an in _type_info}) + + import math + + __deprecated_attrs__['math'] = (math, + "`np.math` is a deprecated alias for the standard library `math` " + "module (Deprecated Numpy 1.25). Replace usages of `np.math` with " + "`math`") + + del math, _msg, _type_info + + from .core import abs + # now that numpy modules are imported, can initialize limits + core.getlimits._register_known_types() + + __all__.extend(['__version__', 'show_config']) + __all__.extend(core.__all__) + __all__.extend(_mat.__all__) + __all__.extend(lib.__all__) + __all__.extend(['linalg', 'fft', 'random', 'ctypeslib', 'ma']) + + # Remove min and max from __all__ to avoid `from numpy import *` override + # the builtins min/max. Temporary fix for 1.25.x/1.26.x, see gh-24229. + __all__.remove('min') + __all__.remove('max') + __all__.remove('round') + + # Remove one of the two occurrences of `issubdtype`, which is exposed as + # both `numpy.core.issubdtype` and `numpy.lib.issubdtype`. + __all__.remove('issubdtype') + + # These are exported by np.core, but are replaced by the builtins below + # remove them to ensure that we don't end up with `np.long == np.int_`, + # which would be a breaking change. + del long, unicode + __all__.remove('long') + __all__.remove('unicode') + + # Remove things that are in the numpy.lib but not in the numpy namespace + # Note that there is a test (numpy/tests/test_public_api.py:test_numpy_namespace) + # that prevents adding more things to the main namespace by accident. + # The list below will grow until the `from .lib import *` fixme above is + # taken care of + __all__.remove('Arrayterator') + del Arrayterator + + # These names were removed in NumPy 1.20. For at least one release, + # attempts to access these names in the numpy namespace will trigger + # a warning, and calling the function will raise an exception. + _financial_names = ['fv', 'ipmt', 'irr', 'mirr', 'nper', 'npv', 'pmt', + 'ppmt', 'pv', 'rate'] + __expired_functions__ = { + name: (f'In accordance with NEP 32, the function {name} was removed ' + 'from NumPy version 1.20. A replacement for this function ' + 'is available in the numpy_financial library: ' + 'https://pypi.org/project/numpy-financial') + for name in _financial_names} + + # Filter out Cython harmless warnings + warnings.filterwarnings("ignore", message="numpy.dtype size changed") + warnings.filterwarnings("ignore", message="numpy.ufunc size changed") + warnings.filterwarnings("ignore", message="numpy.ndarray size changed") + + # oldnumeric and numarray were removed in 1.9. In case some packages import + # but do not use them, we define them here for backward compatibility. + oldnumeric = 'removed' + numarray = 'removed' + + def __getattr__(attr): + # Warn for expired attributes, and return a dummy function + # that always raises an exception. + import warnings + import math + try: + msg = __expired_functions__[attr] + except KeyError: + pass + else: + warnings.warn(msg, DeprecationWarning, stacklevel=2) + + def _expired(*args, **kwds): + raise RuntimeError(msg) + + return _expired + + # Emit warnings for deprecated attributes + try: + val, msg = __deprecated_attrs__[attr] + except KeyError: + pass + else: + warnings.warn(msg, DeprecationWarning, stacklevel=2) + return val + + if attr in __future_scalars__: + # And future warnings for those that will change, but also give + # the AttributeError + warnings.warn( + f"In the future `np.{attr}` will be defined as the " + "corresponding NumPy scalar.", FutureWarning, stacklevel=2) + + if attr in __former_attrs__: + raise AttributeError(__former_attrs__[attr]) + + if attr == 'testing': + import numpy.testing as testing + return testing + elif attr == 'Tester': + "Removed in NumPy 1.25.0" + raise RuntimeError("Tester was removed in NumPy 1.25.") + + raise AttributeError("module {!r} has no attribute " + "{!r}".format(__name__, attr)) + + def __dir__(): + public_symbols = globals().keys() | {'testing'} + public_symbols -= { + "core", "matrixlib", + # These were moved in 1.25 and may be deprecated eventually: + "ModuleDeprecationWarning", "VisibleDeprecationWarning", + "ComplexWarning", "TooHardError", "AxisError" + } + return list(public_symbols) + + # Pytest testing + from numpy._pytesttester import PytestTester + test = PytestTester(__name__) + del PytestTester + + def _sanity_check(): + """ + Quick sanity checks for common bugs caused by environment. + There are some cases e.g. with wrong BLAS ABI that cause wrong + results under specific runtime conditions that are not necessarily + achieved during test suite runs, and it is useful to catch those early. + + See https://github.com/numpy/numpy/issues/8577 and other + similar bug reports. + + """ + try: + x = ones(2, dtype=float32) + if not abs(x.dot(x) - float32(2.0)) < 1e-5: + raise AssertionError() + except AssertionError: + msg = ("The current Numpy installation ({!r}) fails to " + "pass simple sanity checks. This can be caused for example " + "by incorrect BLAS library being linked in, or by mixing " + "package managers (pip, conda, apt, ...). Search closed " + "numpy issues for similar problems.") + raise RuntimeError(msg.format(__file__)) from None + + _sanity_check() + del _sanity_check + + def _mac_os_check(): + """ + Quick Sanity check for Mac OS look for accelerate build bugs. + Testing numpy polyfit calls init_dgelsd(LAPACK) + """ + try: + c = array([3., 2., 1.]) + x = linspace(0, 2, 5) + y = polyval(c, x) + _ = polyfit(x, y, 2, cov=True) + except ValueError: + pass + + if sys.platform == "darwin": + from . import exceptions + with warnings.catch_warnings(record=True) as w: + _mac_os_check() + # Throw runtime error, if the test failed Check for warning and error_message + if len(w) > 0: + for _wn in w: + if _wn.category is exceptions.RankWarning: + # Ignore other warnings, they may not be relevant (see gh-25433). + error_message = f"{_wn.category.__name__}: {str(_wn.message)}" + msg = ( + "Polyfit sanity test emitted a warning, most likely due " + "to using a buggy Accelerate backend." + "\nIf you compiled yourself, more information is available at:" + "\nhttps://numpy.org/devdocs/building/index.html" + "\nOtherwise report this to the vendor " + "that provided NumPy.\n\n{}\n".format(error_message)) + raise RuntimeError(msg) + del _wn + del w + del _mac_os_check + + # We usually use madvise hugepages support, but on some old kernels it + # is slow and thus better avoided. + # Specifically kernel version 4.6 had a bug fix which probably fixed this: + # https://github.com/torvalds/linux/commit/7cf91a98e607c2f935dbcc177d70011e95b8faff + import os + use_hugepage = os.environ.get("NUMPY_MADVISE_HUGEPAGE", None) + if sys.platform == "linux" and use_hugepage is None: + # If there is an issue with parsing the kernel version, + # set use_hugepages to 0. Usage of LooseVersion will handle + # the kernel version parsing better, but avoided since it + # will increase the import time. See: #16679 for related discussion. + try: + use_hugepage = 1 + kernel_version = os.uname().release.split(".")[:2] + kernel_version = tuple(int(v) for v in kernel_version) + if kernel_version < (4, 6): + use_hugepage = 0 + except ValueError: + use_hugepages = 0 + elif use_hugepage is None: + # This is not Linux, so it should not matter, just enable anyway + use_hugepage = 1 + else: + use_hugepage = int(use_hugepage) + + # Note that this will currently only make a difference on Linux + core.multiarray._set_madvise_hugepage(use_hugepage) + del use_hugepage + + # Give a warning if NumPy is reloaded or imported on a sub-interpreter + # We do this from python, since the C-module may not be reloaded and + # it is tidier organized. + core.multiarray._multiarray_umath._reload_guard() + + # default to "weak" promotion for "NumPy 2". + core._set_promotion_state( + os.environ.get("NPY_PROMOTION_STATE", + "weak" if _using_numpy2_behavior() else "legacy")) + + # Tell PyInstaller where to find hook-numpy.py + def _pyinstaller_hooks_dir(): + from pathlib import Path + return [str(Path(__file__).with_name("_pyinstaller").resolve())] + + # Remove symbols imported for internal use + del os + + +# Remove symbols imported for internal use +del sys, warnings diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__init__.pyi b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..a185bfe754e3d6ba5004f8cca06177e60a7aa13c --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/__init__.pyi @@ -0,0 +1,4422 @@ +import builtins +import sys +import os +import mmap +import ctypes as ct +import array as _array +import datetime as dt +import enum +from abc import abstractmethod +from types import TracebackType, MappingProxyType, GenericAlias +from contextlib import ContextDecorator +from contextlib import contextmanager + +from numpy._pytesttester import PytestTester +from numpy.core._internal import _ctypes + +from numpy._typing import ( + # Arrays + ArrayLike, + NDArray, + _SupportsArray, + _NestedSequence, + _FiniteNestedSequence, + _SupportsArray, + _ArrayLikeBool_co, + _ArrayLikeUInt_co, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeNumber_co, + _ArrayLikeTD64_co, + _ArrayLikeDT64_co, + _ArrayLikeObject_co, + _ArrayLikeStr_co, + _ArrayLikeBytes_co, + _ArrayLikeUnknown, + _UnknownType, + + # DTypes + DTypeLike, + _DTypeLike, + _DTypeLikeVoid, + _SupportsDType, + _VoidDTypeLike, + + # Shapes + _Shape, + _ShapeLike, + + # Scalars + _CharLike_co, + _BoolLike_co, + _IntLike_co, + _FloatLike_co, + _ComplexLike_co, + _TD64Like_co, + _NumberLike_co, + _ScalarLike_co, + + # `number` precision + NBitBase, + _256Bit, + _128Bit, + _96Bit, + _80Bit, + _64Bit, + _32Bit, + _16Bit, + _8Bit, + _NBitByte, + _NBitShort, + _NBitIntC, + _NBitIntP, + _NBitInt, + _NBitLongLong, + _NBitHalf, + _NBitSingle, + _NBitDouble, + _NBitLongDouble, + + # Character codes + _BoolCodes, + _UInt8Codes, + _UInt16Codes, + _UInt32Codes, + _UInt64Codes, + _Int8Codes, + _Int16Codes, + _Int32Codes, + _Int64Codes, + _Float16Codes, + _Float32Codes, + _Float64Codes, + _Complex64Codes, + _Complex128Codes, + _ByteCodes, + _ShortCodes, + _IntCCodes, + _IntPCodes, + _IntCodes, + _LongLongCodes, + _UByteCodes, + _UShortCodes, + _UIntCCodes, + _UIntPCodes, + _UIntCodes, + _ULongLongCodes, + _HalfCodes, + _SingleCodes, + _DoubleCodes, + _LongDoubleCodes, + _CSingleCodes, + _CDoubleCodes, + _CLongDoubleCodes, + _DT64Codes, + _TD64Codes, + _StrCodes, + _BytesCodes, + _VoidCodes, + _ObjectCodes, + + # Ufuncs + _UFunc_Nin1_Nout1, + _UFunc_Nin2_Nout1, + _UFunc_Nin1_Nout2, + _UFunc_Nin2_Nout2, + _GUFunc_Nin2_Nout1, +) + +from numpy._typing._callable import ( + _BoolOp, + _BoolBitOp, + _BoolSub, + _BoolTrueDiv, + _BoolMod, + _BoolDivMod, + _TD64Div, + _IntTrueDiv, + _UnsignedIntOp, + _UnsignedIntBitOp, + _UnsignedIntMod, + _UnsignedIntDivMod, + _SignedIntOp, + _SignedIntBitOp, + _SignedIntMod, + _SignedIntDivMod, + _FloatOp, + _FloatMod, + _FloatDivMod, + _ComplexOp, + _NumberOp, + _ComparisonOp, +) + +# NOTE: Numpy's mypy plugin is used for removing the types unavailable +# to the specific platform +from numpy._typing._extended_precision import ( + uint128 as uint128, + uint256 as uint256, + int128 as int128, + int256 as int256, + float80 as float80, + float96 as float96, + float128 as float128, + float256 as float256, + complex160 as complex160, + complex192 as complex192, + complex256 as complex256, + complex512 as complex512, +) + +from collections.abc import ( + Callable, + Container, + Iterable, + Iterator, + Mapping, + Sequence, + Sized, +) +from typing import ( + Literal as L, + Any, + Generator, + Generic, + IO, + NoReturn, + overload, + SupportsComplex, + SupportsFloat, + SupportsInt, + TypeVar, + Union, + Protocol, + SupportsIndex, + Final, + final, + ClassVar, +) + +# Ensures that the stubs are picked up +from numpy import ( + ctypeslib as ctypeslib, + exceptions as exceptions, + fft as fft, + lib as lib, + linalg as linalg, + ma as ma, + polynomial as polynomial, + random as random, + testing as testing, + version as version, + exceptions as exceptions, + dtypes as dtypes, +) + +from numpy.core import defchararray, records +char = defchararray +rec = records + +from numpy.core.function_base import ( + linspace as linspace, + logspace as logspace, + geomspace as geomspace, +) + +from numpy.core.fromnumeric import ( + take as take, + reshape as reshape, + choose as choose, + repeat as repeat, + put as put, + swapaxes as swapaxes, + transpose as transpose, + partition as partition, + argpartition as argpartition, + sort as sort, + argsort as argsort, + argmax as argmax, + argmin as argmin, + searchsorted as searchsorted, + resize as resize, + squeeze as squeeze, + diagonal as diagonal, + trace as trace, + ravel as ravel, + nonzero as nonzero, + shape as shape, + compress as compress, + clip as clip, + sum as sum, + all as all, + any as any, + cumsum as cumsum, + ptp as ptp, + max as max, + min as min, + amax as amax, + amin as amin, + prod as prod, + cumprod as cumprod, + ndim as ndim, + size as size, + around as around, + round as round, + mean as mean, + std as std, + var as var, +) + +from numpy.core._asarray import ( + require as require, +) + +from numpy.core._type_aliases import ( + sctypes as sctypes, + sctypeDict as sctypeDict, +) + +from numpy.core._ufunc_config import ( + seterr as seterr, + geterr as geterr, + setbufsize as setbufsize, + getbufsize as getbufsize, + seterrcall as seterrcall, + geterrcall as geterrcall, + _ErrKind, + _ErrFunc, + _ErrDictOptional, +) + +from numpy.core.arrayprint import ( + set_printoptions as set_printoptions, + get_printoptions as get_printoptions, + array2string as array2string, + format_float_scientific as format_float_scientific, + format_float_positional as format_float_positional, + array_repr as array_repr, + array_str as array_str, + set_string_function as set_string_function, + printoptions as printoptions, +) + +from numpy.core.einsumfunc import ( + einsum as einsum, + einsum_path as einsum_path, +) + +from numpy.core.multiarray import ( + ALLOW_THREADS as ALLOW_THREADS, + BUFSIZE as BUFSIZE, + CLIP as CLIP, + MAXDIMS as MAXDIMS, + MAY_SHARE_BOUNDS as MAY_SHARE_BOUNDS, + MAY_SHARE_EXACT as MAY_SHARE_EXACT, + RAISE as RAISE, + WRAP as WRAP, + tracemalloc_domain as tracemalloc_domain, + array as array, + empty_like as empty_like, + empty as empty, + zeros as zeros, + concatenate as concatenate, + inner as inner, + where as where, + lexsort as lexsort, + can_cast as can_cast, + min_scalar_type as min_scalar_type, + result_type as result_type, + dot as dot, + vdot as vdot, + bincount as bincount, + copyto as copyto, + putmask as putmask, + packbits as packbits, + unpackbits as unpackbits, + shares_memory as shares_memory, + may_share_memory as may_share_memory, + asarray as asarray, + asanyarray as asanyarray, + ascontiguousarray as ascontiguousarray, + asfortranarray as asfortranarray, + arange as arange, + busday_count as busday_count, + busday_offset as busday_offset, + compare_chararrays as compare_chararrays, + datetime_as_string as datetime_as_string, + datetime_data as datetime_data, + frombuffer as frombuffer, + fromfile as fromfile, + fromiter as fromiter, + is_busday as is_busday, + promote_types as promote_types, + seterrobj as seterrobj, + geterrobj as geterrobj, + fromstring as fromstring, + frompyfunc as frompyfunc, + nested_iters as nested_iters, + flagsobj, +) + +from numpy.core.numeric import ( + zeros_like as zeros_like, + ones as ones, + ones_like as ones_like, + full as full, + full_like as full_like, + count_nonzero as count_nonzero, + isfortran as isfortran, + argwhere as argwhere, + flatnonzero as flatnonzero, + correlate as correlate, + convolve as convolve, + outer as outer, + tensordot as tensordot, + roll as roll, + rollaxis as rollaxis, + moveaxis as moveaxis, + cross as cross, + indices as indices, + fromfunction as fromfunction, + isscalar as isscalar, + binary_repr as binary_repr, + base_repr as base_repr, + identity as identity, + allclose as allclose, + isclose as isclose, + array_equal as array_equal, + array_equiv as array_equiv, +) + +from numpy.core.numerictypes import ( + maximum_sctype as maximum_sctype, + issctype as issctype, + obj2sctype as obj2sctype, + issubclass_ as issubclass_, + issubsctype as issubsctype, + issubdtype as issubdtype, + sctype2char as sctype2char, + nbytes as nbytes, + cast as cast, + ScalarType as ScalarType, + typecodes as typecodes, +) + +from numpy.core.shape_base import ( + atleast_1d as atleast_1d, + atleast_2d as atleast_2d, + atleast_3d as atleast_3d, + block as block, + hstack as hstack, + stack as stack, + vstack as vstack, +) + +from numpy.exceptions import ( + ComplexWarning as ComplexWarning, + ModuleDeprecationWarning as ModuleDeprecationWarning, + VisibleDeprecationWarning as VisibleDeprecationWarning, + TooHardError as TooHardError, + DTypePromotionError as DTypePromotionError, + AxisError as AxisError, +) + +from numpy.lib import ( + emath as emath, +) + +from numpy.lib.arraypad import ( + pad as pad, +) + +from numpy.lib.arraysetops import ( + ediff1d as ediff1d, + intersect1d as intersect1d, + setxor1d as setxor1d, + union1d as union1d, + setdiff1d as setdiff1d, + unique as unique, + in1d as in1d, + isin as isin, +) + +from numpy.lib.arrayterator import ( + Arrayterator as Arrayterator, +) + +from numpy.lib.function_base import ( + select as select, + piecewise as piecewise, + trim_zeros as trim_zeros, + copy as copy, + iterable as iterable, + percentile as percentile, + diff as diff, + gradient as gradient, + angle as angle, + unwrap as unwrap, + sort_complex as sort_complex, + disp as disp, + flip as flip, + rot90 as rot90, + extract as extract, + place as place, + asarray_chkfinite as asarray_chkfinite, + average as average, + bincount as bincount, + digitize as digitize, + cov as cov, + corrcoef as corrcoef, + median as median, + sinc as sinc, + hamming as hamming, + hanning as hanning, + bartlett as bartlett, + blackman as blackman, + kaiser as kaiser, + trapz as trapz, + i0 as i0, + add_newdoc as add_newdoc, + add_docstring as add_docstring, + meshgrid as meshgrid, + delete as delete, + insert as insert, + append as append, + interp as interp, + add_newdoc_ufunc as add_newdoc_ufunc, + quantile as quantile, +) + +from numpy.lib.histograms import ( + histogram_bin_edges as histogram_bin_edges, + histogram as histogram, + histogramdd as histogramdd, +) + +from numpy.lib.index_tricks import ( + ravel_multi_index as ravel_multi_index, + unravel_index as unravel_index, + mgrid as mgrid, + ogrid as ogrid, + r_ as r_, + c_ as c_, + s_ as s_, + index_exp as index_exp, + ix_ as ix_, + fill_diagonal as fill_diagonal, + diag_indices as diag_indices, + diag_indices_from as diag_indices_from, +) + +from numpy.lib.nanfunctions import ( + nansum as nansum, + nanmax as nanmax, + nanmin as nanmin, + nanargmax as nanargmax, + nanargmin as nanargmin, + nanmean as nanmean, + nanmedian as nanmedian, + nanpercentile as nanpercentile, + nanvar as nanvar, + nanstd as nanstd, + nanprod as nanprod, + nancumsum as nancumsum, + nancumprod as nancumprod, + nanquantile as nanquantile, +) + +from numpy.lib.npyio import ( + savetxt as savetxt, + loadtxt as loadtxt, + genfromtxt as genfromtxt, + recfromtxt as recfromtxt, + recfromcsv as recfromcsv, + load as load, + save as save, + savez as savez, + savez_compressed as savez_compressed, + packbits as packbits, + unpackbits as unpackbits, + fromregex as fromregex, +) + +from numpy.lib.polynomial import ( + poly as poly, + roots as roots, + polyint as polyint, + polyder as polyder, + polyadd as polyadd, + polysub as polysub, + polymul as polymul, + polydiv as polydiv, + polyval as polyval, + polyfit as polyfit, +) + +from numpy.lib.shape_base import ( + column_stack as column_stack, + row_stack as row_stack, + dstack as dstack, + array_split as array_split, + split as split, + hsplit as hsplit, + vsplit as vsplit, + dsplit as dsplit, + apply_over_axes as apply_over_axes, + expand_dims as expand_dims, + apply_along_axis as apply_along_axis, + kron as kron, + tile as tile, + get_array_wrap as get_array_wrap, + take_along_axis as take_along_axis, + put_along_axis as put_along_axis, +) + +from numpy.lib.stride_tricks import ( + broadcast_to as broadcast_to, + broadcast_arrays as broadcast_arrays, + broadcast_shapes as broadcast_shapes, +) + +from numpy.lib.twodim_base import ( + diag as diag, + diagflat as diagflat, + eye as eye, + fliplr as fliplr, + flipud as flipud, + tri as tri, + triu as triu, + tril as tril, + vander as vander, + histogram2d as histogram2d, + mask_indices as mask_indices, + tril_indices as tril_indices, + tril_indices_from as tril_indices_from, + triu_indices as triu_indices, + triu_indices_from as triu_indices_from, +) + +from numpy.lib.type_check import ( + mintypecode as mintypecode, + asfarray as asfarray, + real as real, + imag as imag, + iscomplex as iscomplex, + isreal as isreal, + iscomplexobj as iscomplexobj, + isrealobj as isrealobj, + nan_to_num as nan_to_num, + real_if_close as real_if_close, + typename as typename, + common_type as common_type, +) + +from numpy.lib.ufunclike import ( + fix as fix, + isposinf as isposinf, + isneginf as isneginf, +) + +from numpy.lib.utils import ( + issubclass_ as issubclass_, + issubsctype as issubsctype, + issubdtype as issubdtype, + deprecate as deprecate, + deprecate_with_doc as deprecate_with_doc, + get_include as get_include, + info as info, + source as source, + who as who, + lookfor as lookfor, + byte_bounds as byte_bounds, + safe_eval as safe_eval, + show_runtime as show_runtime, +) + +from numpy.matrixlib import ( + asmatrix as asmatrix, + mat as mat, + bmat as bmat, +) + +_AnyStr_contra = TypeVar("_AnyStr_contra", str, bytes, contravariant=True) + +# Protocol for representing file-like-objects accepted +# by `ndarray.tofile` and `fromfile` +class _IOProtocol(Protocol): + def flush(self) -> object: ... + def fileno(self) -> int: ... + def tell(self) -> SupportsIndex: ... + def seek(self, offset: int, whence: int, /) -> object: ... + +# NOTE: `seek`, `write` and `flush` are technically only required +# for `readwrite`/`write` modes +class _MemMapIOProtocol(Protocol): + def flush(self) -> object: ... + def fileno(self) -> SupportsIndex: ... + def tell(self) -> int: ... + def seek(self, offset: int, whence: int, /) -> object: ... + def write(self, s: bytes, /) -> object: ... + @property + def read(self) -> object: ... + +class _SupportsWrite(Protocol[_AnyStr_contra]): + def write(self, s: _AnyStr_contra, /) -> object: ... + +__all__: list[str] +__path__: list[str] +__version__: str +test: PytestTester + +# TODO: Move placeholders to their respective module once +# their annotations are properly implemented +# +# Placeholders for classes + +def show_config() -> None: ... + +_NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray[Any, Any]) +_DTypeScalar_co = TypeVar("_DTypeScalar_co", covariant=True, bound=generic) +_ByteOrder = L["S", "<", ">", "=", "|", "L", "B", "N", "I"] + +@final +class dtype(Generic[_DTypeScalar_co]): + names: None | tuple[builtins.str, ...] + def __hash__(self) -> int: ... + # Overload for subclass of generic + @overload + def __new__( + cls, + dtype: type[_DTypeScalar_co], + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[_DTypeScalar_co]: ... + # Overloads for string aliases, Python types, and some assorted + # other special cases. Order is sometimes important because of the + # subtype relationships + # + # bool < int < float < complex < object + # + # so we have to make sure the overloads for the narrowest type is + # first. + # Builtin types + @overload + def __new__(cls, dtype: type[bool], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[bool_]: ... + @overload + def __new__(cls, dtype: type[int], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int_]: ... + @overload + def __new__(cls, dtype: None | type[float], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[float_]: ... + @overload + def __new__(cls, dtype: type[complex], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[complex_]: ... + @overload + def __new__(cls, dtype: type[builtins.str], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[str_]: ... + @overload + def __new__(cls, dtype: type[bytes], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[bytes_]: ... + + # `unsignedinteger` string-based representations and ctypes + @overload + def __new__(cls, dtype: _UInt8Codes | type[ct.c_uint8], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint8]: ... + @overload + def __new__(cls, dtype: _UInt16Codes | type[ct.c_uint16], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint16]: ... + @overload + def __new__(cls, dtype: _UInt32Codes | type[ct.c_uint32], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint32]: ... + @overload + def __new__(cls, dtype: _UInt64Codes | type[ct.c_uint64], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint64]: ... + @overload + def __new__(cls, dtype: _UByteCodes | type[ct.c_ubyte], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[ubyte]: ... + @overload + def __new__(cls, dtype: _UShortCodes | type[ct.c_ushort], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[ushort]: ... + @overload + def __new__(cls, dtype: _UIntCCodes | type[ct.c_uint], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uintc]: ... + + # NOTE: We're assuming here that `uint_ptr_t == size_t`, + # an assumption that does not hold in rare cases (same for `ssize_t`) + @overload + def __new__(cls, dtype: _UIntPCodes | type[ct.c_void_p] | type[ct.c_size_t], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uintp]: ... + @overload + def __new__(cls, dtype: _UIntCodes | type[ct.c_ulong], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint]: ... + @overload + def __new__(cls, dtype: _ULongLongCodes | type[ct.c_ulonglong], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[ulonglong]: ... + + # `signedinteger` string-based representations and ctypes + @overload + def __new__(cls, dtype: _Int8Codes | type[ct.c_int8], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int8]: ... + @overload + def __new__(cls, dtype: _Int16Codes | type[ct.c_int16], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int16]: ... + @overload + def __new__(cls, dtype: _Int32Codes | type[ct.c_int32], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int32]: ... + @overload + def __new__(cls, dtype: _Int64Codes | type[ct.c_int64], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int64]: ... + @overload + def __new__(cls, dtype: _ByteCodes | type[ct.c_byte], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[byte]: ... + @overload + def __new__(cls, dtype: _ShortCodes | type[ct.c_short], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[short]: ... + @overload + def __new__(cls, dtype: _IntCCodes | type[ct.c_int], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[intc]: ... + @overload + def __new__(cls, dtype: _IntPCodes | type[ct.c_ssize_t], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[intp]: ... + @overload + def __new__(cls, dtype: _IntCodes | type[ct.c_long], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int_]: ... + @overload + def __new__(cls, dtype: _LongLongCodes | type[ct.c_longlong], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[longlong]: ... + + # `floating` string-based representations and ctypes + @overload + def __new__(cls, dtype: _Float16Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[float16]: ... + @overload + def __new__(cls, dtype: _Float32Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[float32]: ... + @overload + def __new__(cls, dtype: _Float64Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[float64]: ... + @overload + def __new__(cls, dtype: _HalfCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[half]: ... + @overload + def __new__(cls, dtype: _SingleCodes | type[ct.c_float], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[single]: ... + @overload + def __new__(cls, dtype: _DoubleCodes | type[ct.c_double], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[double]: ... + @overload + def __new__(cls, dtype: _LongDoubleCodes | type[ct.c_longdouble], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[longdouble]: ... + + # `complexfloating` string-based representations + @overload + def __new__(cls, dtype: _Complex64Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[complex64]: ... + @overload + def __new__(cls, dtype: _Complex128Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[complex128]: ... + @overload + def __new__(cls, dtype: _CSingleCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[csingle]: ... + @overload + def __new__(cls, dtype: _CDoubleCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[cdouble]: ... + @overload + def __new__(cls, dtype: _CLongDoubleCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[clongdouble]: ... + + # Miscellaneous string-based representations and ctypes + @overload + def __new__(cls, dtype: _BoolCodes | type[ct.c_bool], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[bool_]: ... + @overload + def __new__(cls, dtype: _TD64Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[timedelta64]: ... + @overload + def __new__(cls, dtype: _DT64Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[datetime64]: ... + @overload + def __new__(cls, dtype: _StrCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[str_]: ... + @overload + def __new__(cls, dtype: _BytesCodes | type[ct.c_char], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[bytes_]: ... + @overload + def __new__(cls, dtype: _VoidCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[void]: ... + @overload + def __new__(cls, dtype: _ObjectCodes | type[ct.py_object[Any]], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[object_]: ... + + # dtype of a dtype is the same dtype + @overload + def __new__( + cls, + dtype: dtype[_DTypeScalar_co], + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[_DTypeScalar_co]: ... + @overload + def __new__( + cls, + dtype: _SupportsDType[dtype[_DTypeScalar_co]], + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[_DTypeScalar_co]: ... + # Handle strings that can't be expressed as literals; i.e. s1, s2, ... + @overload + def __new__( + cls, + dtype: builtins.str, + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[Any]: ... + # Catchall overload for void-likes + @overload + def __new__( + cls, + dtype: _VoidDTypeLike, + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[void]: ... + # Catchall overload for object-likes + @overload + def __new__( + cls, + dtype: type[object], + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[object_]: ... + + def __class_getitem__(self, item: Any) -> GenericAlias: ... + + @overload + def __getitem__(self: dtype[void], key: list[builtins.str]) -> dtype[void]: ... + @overload + def __getitem__(self: dtype[void], key: builtins.str | SupportsIndex) -> dtype[Any]: ... + + # NOTE: In the future 1-based multiplications will also yield `flexible` dtypes + @overload + def __mul__(self: _DType, value: L[1]) -> _DType: ... + @overload + def __mul__(self: _FlexDType, value: SupportsIndex) -> _FlexDType: ... + @overload + def __mul__(self, value: SupportsIndex) -> dtype[void]: ... + + # NOTE: `__rmul__` seems to be broken when used in combination with + # literals as of mypy 0.902. Set the return-type to `dtype[Any]` for + # now for non-flexible dtypes. + @overload + def __rmul__(self: _FlexDType, value: SupportsIndex) -> _FlexDType: ... + @overload + def __rmul__(self, value: SupportsIndex) -> dtype[Any]: ... + + def __gt__(self, other: DTypeLike) -> bool: ... + def __ge__(self, other: DTypeLike) -> bool: ... + def __lt__(self, other: DTypeLike) -> bool: ... + def __le__(self, other: DTypeLike) -> bool: ... + + # Explicitly defined `__eq__` and `__ne__` to get around mypy's + # `strict_equality` option; even though their signatures are + # identical to their `object`-based counterpart + def __eq__(self, other: Any) -> bool: ... + def __ne__(self, other: Any) -> bool: ... + + @property + def alignment(self) -> int: ... + @property + def base(self) -> dtype[Any]: ... + @property + def byteorder(self) -> builtins.str: ... + @property + def char(self) -> builtins.str: ... + @property + def descr(self) -> list[tuple[builtins.str, builtins.str] | tuple[builtins.str, builtins.str, _Shape]]: ... + @property + def fields( + self, + ) -> None | MappingProxyType[builtins.str, tuple[dtype[Any], int] | tuple[dtype[Any], int, Any]]: ... + @property + def flags(self) -> int: ... + @property + def hasobject(self) -> bool: ... + @property + def isbuiltin(self) -> int: ... + @property + def isnative(self) -> bool: ... + @property + def isalignedstruct(self) -> bool: ... + @property + def itemsize(self) -> int: ... + @property + def kind(self) -> builtins.str: ... + @property + def metadata(self) -> None | MappingProxyType[builtins.str, Any]: ... + @property + def name(self) -> builtins.str: ... + @property + def num(self) -> int: ... + @property + def shape(self) -> _Shape: ... + @property + def ndim(self) -> int: ... + @property + def subdtype(self) -> None | tuple[dtype[Any], _Shape]: ... + def newbyteorder(self: _DType, __new_order: _ByteOrder = ...) -> _DType: ... + @property + def str(self) -> builtins.str: ... + @property + def type(self) -> type[_DTypeScalar_co]: ... + +_ArrayLikeInt = Union[ + int, + integer[Any], + Sequence[Union[int, integer[Any]]], + Sequence[Sequence[Any]], # TODO: wait for support for recursive types + ndarray[Any, Any] +] + +_FlatIterSelf = TypeVar("_FlatIterSelf", bound=flatiter[Any]) + +@final +class flatiter(Generic[_NdArraySubClass]): + __hash__: ClassVar[None] + @property + def base(self) -> _NdArraySubClass: ... + @property + def coords(self) -> _Shape: ... + @property + def index(self) -> int: ... + def copy(self) -> _NdArraySubClass: ... + def __iter__(self: _FlatIterSelf) -> _FlatIterSelf: ... + def __next__(self: flatiter[ndarray[Any, dtype[_ScalarType]]]) -> _ScalarType: ... + def __len__(self) -> int: ... + @overload + def __getitem__( + self: flatiter[ndarray[Any, dtype[_ScalarType]]], + key: int | integer[Any] | tuple[int | integer[Any]], + ) -> _ScalarType: ... + @overload + def __getitem__( + self, + key: _ArrayLikeInt | slice | ellipsis | tuple[_ArrayLikeInt | slice | ellipsis], + ) -> _NdArraySubClass: ... + # TODO: `__setitem__` operates via `unsafe` casting rules, and can + # thus accept any type accepted by the relevant underlying `np.generic` + # constructor. + # This means that `value` must in reality be a supertype of `npt.ArrayLike`. + def __setitem__( + self, + key: _ArrayLikeInt | slice | ellipsis | tuple[_ArrayLikeInt | slice | ellipsis], + value: Any, + ) -> None: ... + @overload + def __array__(self: flatiter[ndarray[Any, _DType]], dtype: None = ..., /) -> ndarray[Any, _DType]: ... + @overload + def __array__(self, dtype: _DType, /) -> ndarray[Any, _DType]: ... + +_OrderKACF = L[None, "K", "A", "C", "F"] +_OrderACF = L[None, "A", "C", "F"] +_OrderCF = L[None, "C", "F"] + +_ModeKind = L["raise", "wrap", "clip"] +_PartitionKind = L["introselect"] +_SortKind = L["quicksort", "mergesort", "heapsort", "stable"] +_SortSide = L["left", "right"] + +_ArraySelf = TypeVar("_ArraySelf", bound=_ArrayOrScalarCommon) + +class _ArrayOrScalarCommon: + @property + def T(self: _ArraySelf) -> _ArraySelf: ... + @property + def data(self) -> memoryview: ... + @property + def flags(self) -> flagsobj: ... + @property + def itemsize(self) -> int: ... + @property + def nbytes(self) -> int: ... + def __bool__(self) -> bool: ... + def __bytes__(self) -> bytes: ... + def __str__(self) -> str: ... + def __repr__(self) -> str: ... + def __copy__(self: _ArraySelf) -> _ArraySelf: ... + def __deepcopy__(self: _ArraySelf, memo: None | dict[int, Any], /) -> _ArraySelf: ... + + # TODO: How to deal with the non-commutative nature of `==` and `!=`? + # xref numpy/numpy#17368 + def __eq__(self, other: Any) -> Any: ... + def __ne__(self, other: Any) -> Any: ... + def copy(self: _ArraySelf, order: _OrderKACF = ...) -> _ArraySelf: ... + def dump(self, file: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _SupportsWrite[bytes]) -> None: ... + def dumps(self) -> bytes: ... + def tobytes(self, order: _OrderKACF = ...) -> bytes: ... + # NOTE: `tostring()` is deprecated and therefore excluded + # def tostring(self, order=...): ... + def tofile( + self, + fid: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _IOProtocol, + sep: str = ..., + format: str = ..., + ) -> None: ... + # generics and 0d arrays return builtin scalars + def tolist(self) -> Any: ... + + @property + def __array_interface__(self) -> dict[str, Any]: ... + @property + def __array_priority__(self) -> float: ... + @property + def __array_struct__(self) -> Any: ... # builtins.PyCapsule + def __setstate__(self, state: tuple[ + SupportsIndex, # version + _ShapeLike, # Shape + _DType_co, # DType + bool, # F-continuous + bytes | list[Any], # Data + ], /) -> None: ... + # a `bool_` is returned when `keepdims=True` and `self` is a 0d array + + @overload + def all( + self, + axis: None = ..., + out: None = ..., + keepdims: L[False] = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> bool_: ... + @overload + def all( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def all( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def any( + self, + axis: None = ..., + out: None = ..., + keepdims: L[False] = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> bool_: ... + @overload + def any( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def any( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def argmax( + self, + axis: None = ..., + out: None = ..., + *, + keepdims: L[False] = ..., + ) -> intp: ... + @overload + def argmax( + self, + axis: SupportsIndex = ..., + out: None = ..., + *, + keepdims: bool = ..., + ) -> Any: ... + @overload + def argmax( + self, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + *, + keepdims: bool = ..., + ) -> _NdArraySubClass: ... + + @overload + def argmin( + self, + axis: None = ..., + out: None = ..., + *, + keepdims: L[False] = ..., + ) -> intp: ... + @overload + def argmin( + self, + axis: SupportsIndex = ..., + out: None = ..., + *, + keepdims: bool = ..., + ) -> Any: ... + @overload + def argmin( + self, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + *, + keepdims: bool = ..., + ) -> _NdArraySubClass: ... + + def argsort( + self, + axis: None | SupportsIndex = ..., + kind: None | _SortKind = ..., + order: None | str | Sequence[str] = ..., + ) -> ndarray[Any, Any]: ... + + @overload + def choose( + self, + choices: ArrayLike, + out: None = ..., + mode: _ModeKind = ..., + ) -> ndarray[Any, Any]: ... + @overload + def choose( + self, + choices: ArrayLike, + out: _NdArraySubClass = ..., + mode: _ModeKind = ..., + ) -> _NdArraySubClass: ... + + @overload + def clip( + self, + min: ArrayLike = ..., + max: None | ArrayLike = ..., + out: None = ..., + **kwargs: Any, + ) -> ndarray[Any, Any]: ... + @overload + def clip( + self, + min: None = ..., + max: ArrayLike = ..., + out: None = ..., + **kwargs: Any, + ) -> ndarray[Any, Any]: ... + @overload + def clip( + self, + min: ArrayLike = ..., + max: None | ArrayLike = ..., + out: _NdArraySubClass = ..., + **kwargs: Any, + ) -> _NdArraySubClass: ... + @overload + def clip( + self, + min: None = ..., + max: ArrayLike = ..., + out: _NdArraySubClass = ..., + **kwargs: Any, + ) -> _NdArraySubClass: ... + + @overload + def compress( + self, + a: ArrayLike, + axis: None | SupportsIndex = ..., + out: None = ..., + ) -> ndarray[Any, Any]: ... + @overload + def compress( + self, + a: ArrayLike, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + def conj(self: _ArraySelf) -> _ArraySelf: ... + + def conjugate(self: _ArraySelf) -> _ArraySelf: ... + + @overload + def cumprod( + self, + axis: None | SupportsIndex = ..., + dtype: DTypeLike = ..., + out: None = ..., + ) -> ndarray[Any, Any]: ... + @overload + def cumprod( + self, + axis: None | SupportsIndex = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + @overload + def cumsum( + self, + axis: None | SupportsIndex = ..., + dtype: DTypeLike = ..., + out: None = ..., + ) -> ndarray[Any, Any]: ... + @overload + def cumsum( + self, + axis: None | SupportsIndex = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + @overload + def max( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def max( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def mean( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def mean( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def min( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def min( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + def newbyteorder( + self: _ArraySelf, + __new_order: _ByteOrder = ..., + ) -> _ArraySelf: ... + + @overload + def prod( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def prod( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def ptp( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + ) -> Any: ... + @overload + def ptp( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + ) -> _NdArraySubClass: ... + + @overload + def round( + self: _ArraySelf, + decimals: SupportsIndex = ..., + out: None = ..., + ) -> _ArraySelf: ... + @overload + def round( + self, + decimals: SupportsIndex = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + @overload + def std( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def std( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def sum( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def sum( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def var( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def var( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + +_DType = TypeVar("_DType", bound=dtype[Any]) +_DType_co = TypeVar("_DType_co", covariant=True, bound=dtype[Any]) +_FlexDType = TypeVar("_FlexDType", bound=dtype[flexible]) + +# TODO: Set the `bound` to something more suitable once we +# have proper shape support +_ShapeType = TypeVar("_ShapeType", bound=Any) +_ShapeType2 = TypeVar("_ShapeType2", bound=Any) +_NumberType = TypeVar("_NumberType", bound=number[Any]) + +if sys.version_info >= (3, 12): + from collections.abc import Buffer as _SupportsBuffer +else: + _SupportsBuffer = ( + bytes + | bytearray + | memoryview + | _array.array[Any] + | mmap.mmap + | NDArray[Any] + | generic + ) + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) +_T_contra = TypeVar("_T_contra", contravariant=True) +_2Tuple = tuple[_T, _T] +_CastingKind = L["no", "equiv", "safe", "same_kind", "unsafe"] + +_ArrayUInt_co = NDArray[Union[bool_, unsignedinteger[Any]]] +_ArrayInt_co = NDArray[Union[bool_, integer[Any]]] +_ArrayFloat_co = NDArray[Union[bool_, integer[Any], floating[Any]]] +_ArrayComplex_co = NDArray[Union[bool_, integer[Any], floating[Any], complexfloating[Any, Any]]] +_ArrayNumber_co = NDArray[Union[bool_, number[Any]]] +_ArrayTD64_co = NDArray[Union[bool_, integer[Any], timedelta64]] + +# Introduce an alias for `dtype` to avoid naming conflicts. +_dtype = dtype + +# `builtins.PyCapsule` unfortunately lacks annotations as of the moment; +# use `Any` as a stopgap measure +_PyCapsule = Any + +class _SupportsItem(Protocol[_T_co]): + def item(self, args: Any, /) -> _T_co: ... + +class _SupportsReal(Protocol[_T_co]): + @property + def real(self) -> _T_co: ... + +class _SupportsImag(Protocol[_T_co]): + @property + def imag(self) -> _T_co: ... + +class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]): + __hash__: ClassVar[None] + @property + def base(self) -> None | ndarray[Any, Any]: ... + @property + def ndim(self) -> int: ... + @property + def size(self) -> int: ... + @property + def real( + self: ndarray[_ShapeType, dtype[_SupportsReal[_ScalarType]]], # type: ignore[type-var] + ) -> ndarray[_ShapeType, _dtype[_ScalarType]]: ... + @real.setter + def real(self, value: ArrayLike) -> None: ... + @property + def imag( + self: ndarray[_ShapeType, dtype[_SupportsImag[_ScalarType]]], # type: ignore[type-var] + ) -> ndarray[_ShapeType, _dtype[_ScalarType]]: ... + @imag.setter + def imag(self, value: ArrayLike) -> None: ... + def __new__( + cls: type[_ArraySelf], + shape: _ShapeLike, + dtype: DTypeLike = ..., + buffer: None | _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: None | _ShapeLike = ..., + order: _OrderKACF = ..., + ) -> _ArraySelf: ... + + if sys.version_info >= (3, 12): + def __buffer__(self, flags: int, /) -> memoryview: ... + + def __class_getitem__(self, item: Any) -> GenericAlias: ... + + @overload + def __array__(self, dtype: None = ..., /) -> ndarray[Any, _DType_co]: ... + @overload + def __array__(self, dtype: _DType, /) -> ndarray[Any, _DType]: ... + + def __array_ufunc__( + self, + ufunc: ufunc, + method: L["__call__", "reduce", "reduceat", "accumulate", "outer", "inner"], + *inputs: Any, + **kwargs: Any, + ) -> Any: ... + + def __array_function__( + self, + func: Callable[..., Any], + types: Iterable[type], + args: Iterable[Any], + kwargs: Mapping[str, Any], + ) -> Any: ... + + # NOTE: In practice any object is accepted by `obj`, but as `__array_finalize__` + # is a pseudo-abstract method the type has been narrowed down in order to + # grant subclasses a bit more flexibility + def __array_finalize__(self, obj: None | NDArray[Any], /) -> None: ... + + def __array_wrap__( + self, + array: ndarray[_ShapeType2, _DType], + context: None | tuple[ufunc, tuple[Any, ...], int] = ..., + /, + ) -> ndarray[_ShapeType2, _DType]: ... + + def __array_prepare__( + self, + array: ndarray[_ShapeType2, _DType], + context: None | tuple[ufunc, tuple[Any, ...], int] = ..., + /, + ) -> ndarray[_ShapeType2, _DType]: ... + + @overload + def __getitem__(self, key: ( + NDArray[integer[Any]] + | NDArray[bool_] + | tuple[NDArray[integer[Any]] | NDArray[bool_], ...] + )) -> ndarray[Any, _DType_co]: ... + @overload + def __getitem__(self, key: SupportsIndex | tuple[SupportsIndex, ...]) -> Any: ... + @overload + def __getitem__(self, key: ( + None + | slice + | ellipsis + | SupportsIndex + | _ArrayLikeInt_co + | tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...] + )) -> ndarray[Any, _DType_co]: ... + @overload + def __getitem__(self: NDArray[void], key: str) -> NDArray[Any]: ... + @overload + def __getitem__(self: NDArray[void], key: list[str]) -> ndarray[_ShapeType, _dtype[void]]: ... + + @property + def ctypes(self) -> _ctypes[int]: ... + @property + def shape(self) -> _Shape: ... + @shape.setter + def shape(self, value: _ShapeLike) -> None: ... + @property + def strides(self) -> _Shape: ... + @strides.setter + def strides(self, value: _ShapeLike) -> None: ... + def byteswap(self: _ArraySelf, inplace: bool = ...) -> _ArraySelf: ... + def fill(self, value: Any) -> None: ... + @property + def flat(self: _NdArraySubClass) -> flatiter[_NdArraySubClass]: ... + + # Use the same output type as that of the underlying `generic` + @overload + def item( + self: ndarray[Any, _dtype[_SupportsItem[_T]]], # type: ignore[type-var] + *args: SupportsIndex, + ) -> _T: ... + @overload + def item( + self: ndarray[Any, _dtype[_SupportsItem[_T]]], # type: ignore[type-var] + args: tuple[SupportsIndex, ...], + /, + ) -> _T: ... + + @overload + def itemset(self, value: Any, /) -> None: ... + @overload + def itemset(self, item: _ShapeLike, value: Any, /) -> None: ... + + @overload + def resize(self, new_shape: _ShapeLike, /, *, refcheck: bool = ...) -> None: ... + @overload + def resize(self, *new_shape: SupportsIndex, refcheck: bool = ...) -> None: ... + + def setflags( + self, write: bool = ..., align: bool = ..., uic: bool = ... + ) -> None: ... + + def squeeze( + self, + axis: None | SupportsIndex | tuple[SupportsIndex, ...] = ..., + ) -> ndarray[Any, _DType_co]: ... + + def swapaxes( + self, + axis1: SupportsIndex, + axis2: SupportsIndex, + ) -> ndarray[Any, _DType_co]: ... + + @overload + def transpose(self: _ArraySelf, axes: None | _ShapeLike, /) -> _ArraySelf: ... + @overload + def transpose(self: _ArraySelf, *axes: SupportsIndex) -> _ArraySelf: ... + + def argpartition( + self, + kth: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + kind: _PartitionKind = ..., + order: None | str | Sequence[str] = ..., + ) -> ndarray[Any, _dtype[intp]]: ... + + def diagonal( + self, + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., + ) -> ndarray[Any, _DType_co]: ... + + # 1D + 1D returns a scalar; + # all other with at least 1 non-0D array return an ndarray. + @overload + def dot(self, b: _ScalarLike_co, out: None = ...) -> ndarray[Any, Any]: ... + @overload + def dot(self, b: ArrayLike, out: None = ...) -> Any: ... # type: ignore[misc] + @overload + def dot(self, b: ArrayLike, out: _NdArraySubClass) -> _NdArraySubClass: ... + + # `nonzero()` is deprecated for 0d arrays/generics + def nonzero(self) -> tuple[ndarray[Any, _dtype[intp]], ...]: ... + + def partition( + self, + kth: _ArrayLikeInt_co, + axis: SupportsIndex = ..., + kind: _PartitionKind = ..., + order: None | str | Sequence[str] = ..., + ) -> None: ... + + # `put` is technically available to `generic`, + # but is pointless as `generic`s are immutable + def put( + self, + ind: _ArrayLikeInt_co, + v: ArrayLike, + mode: _ModeKind = ..., + ) -> None: ... + + @overload + def searchsorted( # type: ignore[misc] + self, # >= 1D array + v: _ScalarLike_co, # 0D array-like + side: _SortSide = ..., + sorter: None | _ArrayLikeInt_co = ..., + ) -> intp: ... + @overload + def searchsorted( + self, # >= 1D array + v: ArrayLike, + side: _SortSide = ..., + sorter: None | _ArrayLikeInt_co = ..., + ) -> ndarray[Any, _dtype[intp]]: ... + + def setfield( + self, + val: ArrayLike, + dtype: DTypeLike, + offset: SupportsIndex = ..., + ) -> None: ... + + def sort( + self, + axis: SupportsIndex = ..., + kind: None | _SortKind = ..., + order: None | str | Sequence[str] = ..., + ) -> None: ... + + @overload + def trace( + self, # >= 2D array + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., + dtype: DTypeLike = ..., + out: None = ..., + ) -> Any: ... + @overload + def trace( + self, # >= 2D array + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + @overload + def take( # type: ignore[misc] + self: ndarray[Any, _dtype[_ScalarType]], + indices: _IntLike_co, + axis: None | SupportsIndex = ..., + out: None = ..., + mode: _ModeKind = ..., + ) -> _ScalarType: ... + @overload + def take( # type: ignore[misc] + self, + indices: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + out: None = ..., + mode: _ModeKind = ..., + ) -> ndarray[Any, _DType_co]: ... + @overload + def take( + self, + indices: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + mode: _ModeKind = ..., + ) -> _NdArraySubClass: ... + + def repeat( + self, + repeats: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + ) -> ndarray[Any, _DType_co]: ... + + def flatten( + self, + order: _OrderKACF = ..., + ) -> ndarray[Any, _DType_co]: ... + + def ravel( + self, + order: _OrderKACF = ..., + ) -> ndarray[Any, _DType_co]: ... + + @overload + def reshape( + self, shape: _ShapeLike, /, *, order: _OrderACF = ... + ) -> ndarray[Any, _DType_co]: ... + @overload + def reshape( + self, *shape: SupportsIndex, order: _OrderACF = ... + ) -> ndarray[Any, _DType_co]: ... + + @overload + def astype( + self, + dtype: _DTypeLike[_ScalarType], + order: _OrderKACF = ..., + casting: _CastingKind = ..., + subok: bool = ..., + copy: bool | _CopyMode = ..., + ) -> NDArray[_ScalarType]: ... + @overload + def astype( + self, + dtype: DTypeLike, + order: _OrderKACF = ..., + casting: _CastingKind = ..., + subok: bool = ..., + copy: bool | _CopyMode = ..., + ) -> NDArray[Any]: ... + + @overload + def view(self: _ArraySelf) -> _ArraySelf: ... + @overload + def view(self, type: type[_NdArraySubClass]) -> _NdArraySubClass: ... + @overload + def view(self, dtype: _DTypeLike[_ScalarType]) -> NDArray[_ScalarType]: ... + @overload + def view(self, dtype: DTypeLike) -> NDArray[Any]: ... + @overload + def view( + self, + dtype: DTypeLike, + type: type[_NdArraySubClass], + ) -> _NdArraySubClass: ... + + @overload + def getfield( + self, + dtype: _DTypeLike[_ScalarType], + offset: SupportsIndex = ... + ) -> NDArray[_ScalarType]: ... + @overload + def getfield( + self, + dtype: DTypeLike, + offset: SupportsIndex = ... + ) -> NDArray[Any]: ... + + # Dispatch to the underlying `generic` via protocols + def __int__( + self: ndarray[Any, _dtype[SupportsInt]], # type: ignore[type-var] + ) -> int: ... + + def __float__( + self: ndarray[Any, _dtype[SupportsFloat]], # type: ignore[type-var] + ) -> float: ... + + def __complex__( + self: ndarray[Any, _dtype[SupportsComplex]], # type: ignore[type-var] + ) -> complex: ... + + def __index__( + self: ndarray[Any, _dtype[SupportsIndex]], # type: ignore[type-var] + ) -> int: ... + + def __len__(self) -> int: ... + def __setitem__(self, key, value): ... + def __iter__(self) -> Any: ... + def __contains__(self, key) -> bool: ... + + # The last overload is for catching recursive objects whose + # nesting is too deep. + # The first overload is for catching `bytes` (as they are a subtype of + # `Sequence[int]`) and `str`. As `str` is a recursive sequence of + # strings, it will pass through the final overload otherwise + + @overload + def __lt__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co) -> NDArray[bool_]: ... + @overload + def __lt__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[bool_]: ... + @overload + def __lt__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[bool_]: ... + @overload + def __lt__(self: NDArray[object_], other: Any) -> NDArray[bool_]: ... + @overload + def __lt__(self: NDArray[Any], other: _ArrayLikeObject_co) -> NDArray[bool_]: ... + + @overload + def __le__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co) -> NDArray[bool_]: ... + @overload + def __le__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[bool_]: ... + @overload + def __le__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[bool_]: ... + @overload + def __le__(self: NDArray[object_], other: Any) -> NDArray[bool_]: ... + @overload + def __le__(self: NDArray[Any], other: _ArrayLikeObject_co) -> NDArray[bool_]: ... + + @overload + def __gt__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co) -> NDArray[bool_]: ... + @overload + def __gt__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[bool_]: ... + @overload + def __gt__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[bool_]: ... + @overload + def __gt__(self: NDArray[object_], other: Any) -> NDArray[bool_]: ... + @overload + def __gt__(self: NDArray[Any], other: _ArrayLikeObject_co) -> NDArray[bool_]: ... + + @overload + def __ge__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co) -> NDArray[bool_]: ... + @overload + def __ge__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[bool_]: ... + @overload + def __ge__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[bool_]: ... + @overload + def __ge__(self: NDArray[object_], other: Any) -> NDArray[bool_]: ... + @overload + def __ge__(self: NDArray[Any], other: _ArrayLikeObject_co) -> NDArray[bool_]: ... + + # Unary ops + @overload + def __abs__(self: NDArray[bool_]) -> NDArray[bool_]: ... + @overload + def __abs__(self: NDArray[complexfloating[_NBit1, _NBit1]]) -> NDArray[floating[_NBit1]]: ... + @overload + def __abs__(self: NDArray[_NumberType]) -> NDArray[_NumberType]: ... + @overload + def __abs__(self: NDArray[timedelta64]) -> NDArray[timedelta64]: ... + @overload + def __abs__(self: NDArray[object_]) -> Any: ... + + @overload + def __invert__(self: NDArray[bool_]) -> NDArray[bool_]: ... + @overload + def __invert__(self: NDArray[_IntType]) -> NDArray[_IntType]: ... + @overload + def __invert__(self: NDArray[object_]) -> Any: ... + + @overload + def __pos__(self: NDArray[_NumberType]) -> NDArray[_NumberType]: ... + @overload + def __pos__(self: NDArray[timedelta64]) -> NDArray[timedelta64]: ... + @overload + def __pos__(self: NDArray[object_]) -> Any: ... + + @overload + def __neg__(self: NDArray[_NumberType]) -> NDArray[_NumberType]: ... + @overload + def __neg__(self: NDArray[timedelta64]) -> NDArray[timedelta64]: ... + @overload + def __neg__(self: NDArray[object_]) -> Any: ... + + # Binary ops + @overload + def __matmul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __matmul__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __matmul__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __matmul__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __matmul__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + @overload + def __matmul__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __matmul__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __matmul__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rmatmul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __rmatmul__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmatmul__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmatmul__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rmatmul__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + @overload + def __rmatmul__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rmatmul__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rmatmul__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __mod__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __mod__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __mod__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __mod__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __mod__(self: _ArrayTD64_co, other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[timedelta64]: ... + @overload + def __mod__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __mod__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rmod__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rmod__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmod__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmod__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rmod__(self: _ArrayTD64_co, other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[timedelta64]: ... + @overload + def __rmod__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rmod__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __divmod__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> _2Tuple[NDArray[int8]]: ... # type: ignore[misc] + @overload + def __divmod__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> _2Tuple[NDArray[unsignedinteger[Any]]]: ... # type: ignore[misc] + @overload + def __divmod__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> _2Tuple[NDArray[signedinteger[Any]]]: ... # type: ignore[misc] + @overload + def __divmod__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> _2Tuple[NDArray[floating[Any]]]: ... # type: ignore[misc] + @overload + def __divmod__(self: _ArrayTD64_co, other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> tuple[NDArray[int64], NDArray[timedelta64]]: ... + + @overload + def __rdivmod__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> _2Tuple[NDArray[int8]]: ... # type: ignore[misc] + @overload + def __rdivmod__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> _2Tuple[NDArray[unsignedinteger[Any]]]: ... # type: ignore[misc] + @overload + def __rdivmod__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> _2Tuple[NDArray[signedinteger[Any]]]: ... # type: ignore[misc] + @overload + def __rdivmod__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> _2Tuple[NDArray[floating[Any]]]: ... # type: ignore[misc] + @overload + def __rdivmod__(self: _ArrayTD64_co, other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> tuple[NDArray[int64], NDArray[timedelta64]]: ... + + @overload + def __add__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __add__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __add__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayTD64_co, other: _ArrayLikeDT64_co) -> NDArray[datetime64]: ... + @overload + def __add__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __add__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __add__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __radd__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __radd__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __radd__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayTD64_co, other: _ArrayLikeDT64_co) -> NDArray[datetime64]: ... + @overload + def __radd__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __radd__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __radd__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __sub__(self: NDArray[_UnknownType], other: _ArrayLikeUnknown) -> NDArray[Any]: ... + @overload + def __sub__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __sub__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __sub__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __sub__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __sub__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __sub__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __sub__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... # type: ignore[misc] + @overload + def __sub__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __sub__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[timedelta64]: ... + @overload + def __sub__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __sub__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rsub__(self: NDArray[_UnknownType], other: _ArrayLikeUnknown) -> NDArray[Any]: ... + @overload + def __rsub__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __rsub__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rsub__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rsub__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rsub__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __rsub__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rsub__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... # type: ignore[misc] + @overload + def __rsub__(self: _ArrayTD64_co, other: _ArrayLikeDT64_co) -> NDArray[datetime64]: ... # type: ignore[misc] + @overload + def __rsub__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[timedelta64]: ... + @overload + def __rsub__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rsub__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __mul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __mul__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __mul__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __mul__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __mul__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __mul__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __mul__(self: _ArrayTD64_co, other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __mul__(self: _ArrayFloat_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __mul__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __mul__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rmul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __rmul__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmul__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmul__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rmul__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __rmul__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rmul__(self: _ArrayTD64_co, other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __rmul__(self: _ArrayFloat_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __rmul__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rmul__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __floordiv__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __floordiv__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __floordiv__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __floordiv__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __floordiv__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[int64]: ... + @overload + def __floordiv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __floordiv__(self: NDArray[timedelta64], other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __floordiv__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __floordiv__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rfloordiv__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rfloordiv__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rfloordiv__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rfloordiv__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rfloordiv__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[int64]: ... + @overload + def __rfloordiv__(self: NDArray[bool_], other: _ArrayLikeTD64_co) -> NoReturn: ... + @overload + def __rfloordiv__(self: _ArrayFloat_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __rfloordiv__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rfloordiv__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __pow__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __pow__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __pow__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __pow__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __pow__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + @overload + def __pow__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __pow__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __pow__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rpow__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rpow__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rpow__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rpow__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rpow__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + @overload + def __rpow__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rpow__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rpow__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __truediv__(self: _ArrayInt_co, other: _ArrayInt_co) -> NDArray[float64]: ... # type: ignore[misc] + @overload + def __truediv__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __truediv__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __truediv__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __truediv__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[float64]: ... + @overload + def __truediv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __truediv__(self: NDArray[timedelta64], other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __truediv__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __truediv__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rtruediv__(self: _ArrayInt_co, other: _ArrayInt_co) -> NDArray[float64]: ... # type: ignore[misc] + @overload + def __rtruediv__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rtruediv__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __rtruediv__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rtruediv__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[float64]: ... + @overload + def __rtruediv__(self: NDArray[bool_], other: _ArrayLikeTD64_co) -> NoReturn: ... + @overload + def __rtruediv__(self: _ArrayFloat_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __rtruediv__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rtruediv__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __lshift__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __lshift__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __lshift__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __lshift__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __lshift__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rlshift__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rlshift__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rlshift__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rlshift__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rlshift__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rshift__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rshift__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rshift__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rshift__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rshift__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rrshift__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rrshift__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rrshift__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rrshift__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rrshift__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __and__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __and__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __and__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __and__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __and__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rand__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __rand__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rand__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rand__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rand__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __xor__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __xor__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __xor__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __xor__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __xor__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rxor__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __rxor__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rxor__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rxor__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rxor__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __or__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __or__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __or__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __or__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __or__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __ror__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __ror__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __ror__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __ror__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __ror__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + # `np.generic` does not support inplace operations + + # NOTE: Inplace ops generally use "same_kind" casting w.r.t. to the left + # operand. An exception to this rule are unsigned integers though, which + # also accepts a signed integer for the right operand as long it is a 0D + # object and its value is >= 0 + @overload + def __iadd__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __iadd__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __iadd__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __iadd__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __iadd__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __iadd__(self: NDArray[timedelta64], other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __iadd__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __iadd__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __isub__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __isub__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __isub__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __isub__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __isub__(self: NDArray[timedelta64], other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __isub__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __isub__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __imul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __imul__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __imul__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __imul__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __imul__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __imul__(self: NDArray[timedelta64], other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __imul__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __itruediv__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __itruediv__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __itruediv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __itruediv__(self: NDArray[timedelta64], other: _ArrayLikeInt_co) -> NDArray[timedelta64]: ... + @overload + def __itruediv__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ifloordiv__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ifloordiv__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ifloordiv__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __ifloordiv__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __ifloordiv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __ifloordiv__(self: NDArray[timedelta64], other: _ArrayLikeInt_co) -> NDArray[timedelta64]: ... + @overload + def __ifloordiv__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ipow__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ipow__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ipow__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __ipow__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __ipow__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __imod__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __imod__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __imod__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __imod__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[timedelta64]: ... + @overload + def __imod__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ilshift__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ilshift__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ilshift__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __irshift__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __irshift__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __irshift__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __iand__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __iand__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __iand__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __iand__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ixor__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __ixor__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ixor__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ixor__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ior__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __ior__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ior__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ior__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __imatmul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __imatmul__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __imatmul__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __imatmul__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __imatmul__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __imatmul__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + def __dlpack__(self: NDArray[number[Any]], *, stream: None = ...) -> _PyCapsule: ... + def __dlpack_device__(self) -> tuple[int, L[0]]: ... + + # Keep `dtype` at the bottom to avoid name conflicts with `np.dtype` + @property + def dtype(self) -> _DType_co: ... + +# NOTE: while `np.generic` is not technically an instance of `ABCMeta`, +# the `@abstractmethod` decorator is herein used to (forcefully) deny +# the creation of `np.generic` instances. +# The `# type: ignore` comments are necessary to silence mypy errors regarding +# the missing `ABCMeta` metaclass. + +# See https://github.com/numpy/numpy-stubs/pull/80 for more details. + +_ScalarType = TypeVar("_ScalarType", bound=generic) +_NBit1 = TypeVar("_NBit1", bound=NBitBase) +_NBit2 = TypeVar("_NBit2", bound=NBitBase) + +class generic(_ArrayOrScalarCommon): + @abstractmethod + def __init__(self, *args: Any, **kwargs: Any) -> None: ... + @overload + def __array__(self: _ScalarType, dtype: None = ..., /) -> ndarray[Any, _dtype[_ScalarType]]: ... + @overload + def __array__(self, dtype: _DType, /) -> ndarray[Any, _DType]: ... + def __hash__(self) -> int: ... + @property + def base(self) -> None: ... + @property + def ndim(self) -> L[0]: ... + @property + def size(self) -> L[1]: ... + @property + def shape(self) -> tuple[()]: ... + @property + def strides(self) -> tuple[()]: ... + def byteswap(self: _ScalarType, inplace: L[False] = ...) -> _ScalarType: ... + @property + def flat(self: _ScalarType) -> flatiter[ndarray[Any, _dtype[_ScalarType]]]: ... + + if sys.version_info >= (3, 12): + def __buffer__(self, flags: int, /) -> memoryview: ... + + @overload + def astype( + self, + dtype: _DTypeLike[_ScalarType], + order: _OrderKACF = ..., + casting: _CastingKind = ..., + subok: bool = ..., + copy: bool | _CopyMode = ..., + ) -> _ScalarType: ... + @overload + def astype( + self, + dtype: DTypeLike, + order: _OrderKACF = ..., + casting: _CastingKind = ..., + subok: bool = ..., + copy: bool | _CopyMode = ..., + ) -> Any: ... + + # NOTE: `view` will perform a 0D->scalar cast, + # thus the array `type` is irrelevant to the output type + @overload + def view( + self: _ScalarType, + type: type[ndarray[Any, Any]] = ..., + ) -> _ScalarType: ... + @overload + def view( + self, + dtype: _DTypeLike[_ScalarType], + type: type[ndarray[Any, Any]] = ..., + ) -> _ScalarType: ... + @overload + def view( + self, + dtype: DTypeLike, + type: type[ndarray[Any, Any]] = ..., + ) -> Any: ... + + @overload + def getfield( + self, + dtype: _DTypeLike[_ScalarType], + offset: SupportsIndex = ... + ) -> _ScalarType: ... + @overload + def getfield( + self, + dtype: DTypeLike, + offset: SupportsIndex = ... + ) -> Any: ... + + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> Any: ... + + @overload + def take( # type: ignore[misc] + self: _ScalarType, + indices: _IntLike_co, + axis: None | SupportsIndex = ..., + out: None = ..., + mode: _ModeKind = ..., + ) -> _ScalarType: ... + @overload + def take( # type: ignore[misc] + self: _ScalarType, + indices: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + out: None = ..., + mode: _ModeKind = ..., + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + @overload + def take( + self, + indices: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + mode: _ModeKind = ..., + ) -> _NdArraySubClass: ... + + def repeat( + self: _ScalarType, + repeats: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + + def flatten( + self: _ScalarType, + order: _OrderKACF = ..., + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + + def ravel( + self: _ScalarType, + order: _OrderKACF = ..., + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + + @overload + def reshape( + self: _ScalarType, shape: _ShapeLike, /, *, order: _OrderACF = ... + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + @overload + def reshape( + self: _ScalarType, *shape: SupportsIndex, order: _OrderACF = ... + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + + def squeeze( + self: _ScalarType, axis: None | L[0] | tuple[()] = ... + ) -> _ScalarType: ... + def transpose(self: _ScalarType, axes: None | tuple[()] = ..., /) -> _ScalarType: ... + # Keep `dtype` at the bottom to avoid name conflicts with `np.dtype` + @property + def dtype(self: _ScalarType) -> _dtype[_ScalarType]: ... + +class number(generic, Generic[_NBit1]): # type: ignore + @property + def real(self: _ArraySelf) -> _ArraySelf: ... + @property + def imag(self: _ArraySelf) -> _ArraySelf: ... + def __class_getitem__(self, item: Any) -> GenericAlias: ... + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def __complex__(self) -> complex: ... + def __neg__(self: _ArraySelf) -> _ArraySelf: ... + def __pos__(self: _ArraySelf) -> _ArraySelf: ... + def __abs__(self: _ArraySelf) -> _ArraySelf: ... + # Ensure that objects annotated as `number` support arithmetic operations + __add__: _NumberOp + __radd__: _NumberOp + __sub__: _NumberOp + __rsub__: _NumberOp + __mul__: _NumberOp + __rmul__: _NumberOp + __floordiv__: _NumberOp + __rfloordiv__: _NumberOp + __pow__: _NumberOp + __rpow__: _NumberOp + __truediv__: _NumberOp + __rtruediv__: _NumberOp + __lt__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __le__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __gt__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __ge__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + +class bool_(generic): + def __init__(self, value: object = ..., /) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> bool: ... + def tolist(self) -> bool: ... + @property + def real(self: _ArraySelf) -> _ArraySelf: ... + @property + def imag(self: _ArraySelf) -> _ArraySelf: ... + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def __complex__(self) -> complex: ... + def __abs__(self: _ArraySelf) -> _ArraySelf: ... + __add__: _BoolOp[bool_] + __radd__: _BoolOp[bool_] + __sub__: _BoolSub + __rsub__: _BoolSub + __mul__: _BoolOp[bool_] + __rmul__: _BoolOp[bool_] + __floordiv__: _BoolOp[int8] + __rfloordiv__: _BoolOp[int8] + __pow__: _BoolOp[int8] + __rpow__: _BoolOp[int8] + __truediv__: _BoolTrueDiv + __rtruediv__: _BoolTrueDiv + def __invert__(self) -> bool_: ... + __lshift__: _BoolBitOp[int8] + __rlshift__: _BoolBitOp[int8] + __rshift__: _BoolBitOp[int8] + __rrshift__: _BoolBitOp[int8] + __and__: _BoolBitOp[bool_] + __rand__: _BoolBitOp[bool_] + __xor__: _BoolBitOp[bool_] + __rxor__: _BoolBitOp[bool_] + __or__: _BoolBitOp[bool_] + __ror__: _BoolBitOp[bool_] + __mod__: _BoolMod + __rmod__: _BoolMod + __divmod__: _BoolDivMod + __rdivmod__: _BoolDivMod + __lt__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __le__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __gt__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __ge__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + +class object_(generic): + def __init__(self, value: object = ..., /) -> None: ... + @property + def real(self: _ArraySelf) -> _ArraySelf: ... + @property + def imag(self: _ArraySelf) -> _ArraySelf: ... + # The 3 protocols below may or may not raise, + # depending on the underlying object + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def __complex__(self) -> complex: ... + + if sys.version_info >= (3, 12): + def __release_buffer__(self, buffer: memoryview, /) -> None: ... + +# The `datetime64` constructors requires an object with the three attributes below, +# and thus supports datetime duck typing +class _DatetimeScalar(Protocol): + @property + def day(self) -> int: ... + @property + def month(self) -> int: ... + @property + def year(self) -> int: ... + +# TODO: `item`/`tolist` returns either `dt.date`, `dt.datetime` or `int` +# depending on the unit +class datetime64(generic): + @overload + def __init__( + self, + value: None | datetime64 | _CharLike_co | _DatetimeScalar = ..., + format: _CharLike_co | tuple[_CharLike_co, _IntLike_co] = ..., + /, + ) -> None: ... + @overload + def __init__( + self, + value: int, + format: _CharLike_co | tuple[_CharLike_co, _IntLike_co], + /, + ) -> None: ... + def __add__(self, other: _TD64Like_co) -> datetime64: ... + def __radd__(self, other: _TD64Like_co) -> datetime64: ... + @overload + def __sub__(self, other: datetime64) -> timedelta64: ... + @overload + def __sub__(self, other: _TD64Like_co) -> datetime64: ... + def __rsub__(self, other: datetime64) -> timedelta64: ... + __lt__: _ComparisonOp[datetime64, _ArrayLikeDT64_co] + __le__: _ComparisonOp[datetime64, _ArrayLikeDT64_co] + __gt__: _ComparisonOp[datetime64, _ArrayLikeDT64_co] + __ge__: _ComparisonOp[datetime64, _ArrayLikeDT64_co] + +_IntValue = Union[SupportsInt, _CharLike_co, SupportsIndex] +_FloatValue = Union[None, _CharLike_co, SupportsFloat, SupportsIndex] +_ComplexValue = Union[ + None, + _CharLike_co, + SupportsFloat, + SupportsComplex, + SupportsIndex, + complex, # `complex` is not a subtype of `SupportsComplex` +] + +class integer(number[_NBit1]): # type: ignore + @property + def numerator(self: _ScalarType) -> _ScalarType: ... + @property + def denominator(self) -> L[1]: ... + @overload + def __round__(self, ndigits: None = ...) -> int: ... + @overload + def __round__(self: _ScalarType, ndigits: SupportsIndex) -> _ScalarType: ... + + # NOTE: `__index__` is technically defined in the bottom-most + # sub-classes (`int64`, `uint32`, etc) + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> int: ... + def tolist(self) -> int: ... + def is_integer(self) -> L[True]: ... + def bit_count(self: _ScalarType) -> int: ... + def __index__(self) -> int: ... + __truediv__: _IntTrueDiv[_NBit1] + __rtruediv__: _IntTrueDiv[_NBit1] + def __mod__(self, value: _IntLike_co) -> integer[Any]: ... + def __rmod__(self, value: _IntLike_co) -> integer[Any]: ... + def __invert__(self: _IntType) -> _IntType: ... + # Ensure that objects annotated as `integer` support bit-wise operations + def __lshift__(self, other: _IntLike_co) -> integer[Any]: ... + def __rlshift__(self, other: _IntLike_co) -> integer[Any]: ... + def __rshift__(self, other: _IntLike_co) -> integer[Any]: ... + def __rrshift__(self, other: _IntLike_co) -> integer[Any]: ... + def __and__(self, other: _IntLike_co) -> integer[Any]: ... + def __rand__(self, other: _IntLike_co) -> integer[Any]: ... + def __or__(self, other: _IntLike_co) -> integer[Any]: ... + def __ror__(self, other: _IntLike_co) -> integer[Any]: ... + def __xor__(self, other: _IntLike_co) -> integer[Any]: ... + def __rxor__(self, other: _IntLike_co) -> integer[Any]: ... + +class signedinteger(integer[_NBit1]): + def __init__(self, value: _IntValue = ..., /) -> None: ... + __add__: _SignedIntOp[_NBit1] + __radd__: _SignedIntOp[_NBit1] + __sub__: _SignedIntOp[_NBit1] + __rsub__: _SignedIntOp[_NBit1] + __mul__: _SignedIntOp[_NBit1] + __rmul__: _SignedIntOp[_NBit1] + __floordiv__: _SignedIntOp[_NBit1] + __rfloordiv__: _SignedIntOp[_NBit1] + __pow__: _SignedIntOp[_NBit1] + __rpow__: _SignedIntOp[_NBit1] + __lshift__: _SignedIntBitOp[_NBit1] + __rlshift__: _SignedIntBitOp[_NBit1] + __rshift__: _SignedIntBitOp[_NBit1] + __rrshift__: _SignedIntBitOp[_NBit1] + __and__: _SignedIntBitOp[_NBit1] + __rand__: _SignedIntBitOp[_NBit1] + __xor__: _SignedIntBitOp[_NBit1] + __rxor__: _SignedIntBitOp[_NBit1] + __or__: _SignedIntBitOp[_NBit1] + __ror__: _SignedIntBitOp[_NBit1] + __mod__: _SignedIntMod[_NBit1] + __rmod__: _SignedIntMod[_NBit1] + __divmod__: _SignedIntDivMod[_NBit1] + __rdivmod__: _SignedIntDivMod[_NBit1] + +int8 = signedinteger[_8Bit] +int16 = signedinteger[_16Bit] +int32 = signedinteger[_32Bit] +int64 = signedinteger[_64Bit] + +byte = signedinteger[_NBitByte] +short = signedinteger[_NBitShort] +intc = signedinteger[_NBitIntC] +intp = signedinteger[_NBitIntP] +int_ = signedinteger[_NBitInt] +longlong = signedinteger[_NBitLongLong] + +# TODO: `item`/`tolist` returns either `dt.timedelta` or `int` +# depending on the unit +class timedelta64(generic): + def __init__( + self, + value: None | int | _CharLike_co | dt.timedelta | timedelta64 = ..., + format: _CharLike_co | tuple[_CharLike_co, _IntLike_co] = ..., + /, + ) -> None: ... + @property + def numerator(self: _ScalarType) -> _ScalarType: ... + @property + def denominator(self) -> L[1]: ... + + # NOTE: Only a limited number of units support conversion + # to builtin scalar types: `Y`, `M`, `ns`, `ps`, `fs`, `as` + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def __complex__(self) -> complex: ... + def __neg__(self: _ArraySelf) -> _ArraySelf: ... + def __pos__(self: _ArraySelf) -> _ArraySelf: ... + def __abs__(self: _ArraySelf) -> _ArraySelf: ... + def __add__(self, other: _TD64Like_co) -> timedelta64: ... + def __radd__(self, other: _TD64Like_co) -> timedelta64: ... + def __sub__(self, other: _TD64Like_co) -> timedelta64: ... + def __rsub__(self, other: _TD64Like_co) -> timedelta64: ... + def __mul__(self, other: _FloatLike_co) -> timedelta64: ... + def __rmul__(self, other: _FloatLike_co) -> timedelta64: ... + __truediv__: _TD64Div[float64] + __floordiv__: _TD64Div[int64] + def __rtruediv__(self, other: timedelta64) -> float64: ... + def __rfloordiv__(self, other: timedelta64) -> int64: ... + def __mod__(self, other: timedelta64) -> timedelta64: ... + def __rmod__(self, other: timedelta64) -> timedelta64: ... + def __divmod__(self, other: timedelta64) -> tuple[int64, timedelta64]: ... + def __rdivmod__(self, other: timedelta64) -> tuple[int64, timedelta64]: ... + __lt__: _ComparisonOp[_TD64Like_co, _ArrayLikeTD64_co] + __le__: _ComparisonOp[_TD64Like_co, _ArrayLikeTD64_co] + __gt__: _ComparisonOp[_TD64Like_co, _ArrayLikeTD64_co] + __ge__: _ComparisonOp[_TD64Like_co, _ArrayLikeTD64_co] + +class unsignedinteger(integer[_NBit1]): + # NOTE: `uint64 + signedinteger -> float64` + def __init__(self, value: _IntValue = ..., /) -> None: ... + __add__: _UnsignedIntOp[_NBit1] + __radd__: _UnsignedIntOp[_NBit1] + __sub__: _UnsignedIntOp[_NBit1] + __rsub__: _UnsignedIntOp[_NBit1] + __mul__: _UnsignedIntOp[_NBit1] + __rmul__: _UnsignedIntOp[_NBit1] + __floordiv__: _UnsignedIntOp[_NBit1] + __rfloordiv__: _UnsignedIntOp[_NBit1] + __pow__: _UnsignedIntOp[_NBit1] + __rpow__: _UnsignedIntOp[_NBit1] + __lshift__: _UnsignedIntBitOp[_NBit1] + __rlshift__: _UnsignedIntBitOp[_NBit1] + __rshift__: _UnsignedIntBitOp[_NBit1] + __rrshift__: _UnsignedIntBitOp[_NBit1] + __and__: _UnsignedIntBitOp[_NBit1] + __rand__: _UnsignedIntBitOp[_NBit1] + __xor__: _UnsignedIntBitOp[_NBit1] + __rxor__: _UnsignedIntBitOp[_NBit1] + __or__: _UnsignedIntBitOp[_NBit1] + __ror__: _UnsignedIntBitOp[_NBit1] + __mod__: _UnsignedIntMod[_NBit1] + __rmod__: _UnsignedIntMod[_NBit1] + __divmod__: _UnsignedIntDivMod[_NBit1] + __rdivmod__: _UnsignedIntDivMod[_NBit1] + +uint8 = unsignedinteger[_8Bit] +uint16 = unsignedinteger[_16Bit] +uint32 = unsignedinteger[_32Bit] +uint64 = unsignedinteger[_64Bit] + +ubyte = unsignedinteger[_NBitByte] +ushort = unsignedinteger[_NBitShort] +uintc = unsignedinteger[_NBitIntC] +uintp = unsignedinteger[_NBitIntP] +uint = unsignedinteger[_NBitInt] +ulonglong = unsignedinteger[_NBitLongLong] + +class inexact(number[_NBit1]): # type: ignore + def __getnewargs__(self: inexact[_64Bit]) -> tuple[float, ...]: ... + +_IntType = TypeVar("_IntType", bound=integer[Any]) +_FloatType = TypeVar('_FloatType', bound=floating[Any]) + +class floating(inexact[_NBit1]): + def __init__(self, value: _FloatValue = ..., /) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., + /, + ) -> float: ... + def tolist(self) -> float: ... + def is_integer(self) -> bool: ... + def hex(self: float64) -> str: ... + @classmethod + def fromhex(cls: type[float64], string: str, /) -> float64: ... + def as_integer_ratio(self) -> tuple[int, int]: ... + def __ceil__(self: float64) -> int: ... + def __floor__(self: float64) -> int: ... + def __trunc__(self: float64) -> int: ... + def __getnewargs__(self: float64) -> tuple[float]: ... + def __getformat__(self: float64, typestr: L["double", "float"], /) -> str: ... + @overload + def __round__(self, ndigits: None = ...) -> int: ... + @overload + def __round__(self: _ScalarType, ndigits: SupportsIndex) -> _ScalarType: ... + __add__: _FloatOp[_NBit1] + __radd__: _FloatOp[_NBit1] + __sub__: _FloatOp[_NBit1] + __rsub__: _FloatOp[_NBit1] + __mul__: _FloatOp[_NBit1] + __rmul__: _FloatOp[_NBit1] + __truediv__: _FloatOp[_NBit1] + __rtruediv__: _FloatOp[_NBit1] + __floordiv__: _FloatOp[_NBit1] + __rfloordiv__: _FloatOp[_NBit1] + __pow__: _FloatOp[_NBit1] + __rpow__: _FloatOp[_NBit1] + __mod__: _FloatMod[_NBit1] + __rmod__: _FloatMod[_NBit1] + __divmod__: _FloatDivMod[_NBit1] + __rdivmod__: _FloatDivMod[_NBit1] + +float16 = floating[_16Bit] +float32 = floating[_32Bit] +float64 = floating[_64Bit] + +half = floating[_NBitHalf] +single = floating[_NBitSingle] +double = floating[_NBitDouble] +float_ = floating[_NBitDouble] +longdouble = floating[_NBitLongDouble] +longfloat = floating[_NBitLongDouble] + +# The main reason for `complexfloating` having two typevars is cosmetic. +# It is used to clarify why `complex128`s precision is `_64Bit`, the latter +# describing the two 64 bit floats representing its real and imaginary component + +class complexfloating(inexact[_NBit1], Generic[_NBit1, _NBit2]): + def __init__(self, value: _ComplexValue = ..., /) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> complex: ... + def tolist(self) -> complex: ... + @property + def real(self) -> floating[_NBit1]: ... # type: ignore[override] + @property + def imag(self) -> floating[_NBit2]: ... # type: ignore[override] + def __abs__(self) -> floating[_NBit1]: ... # type: ignore[override] + def __getnewargs__(self: complex128) -> tuple[float, float]: ... + # NOTE: Deprecated + # def __round__(self, ndigits=...): ... + __add__: _ComplexOp[_NBit1] + __radd__: _ComplexOp[_NBit1] + __sub__: _ComplexOp[_NBit1] + __rsub__: _ComplexOp[_NBit1] + __mul__: _ComplexOp[_NBit1] + __rmul__: _ComplexOp[_NBit1] + __truediv__: _ComplexOp[_NBit1] + __rtruediv__: _ComplexOp[_NBit1] + __pow__: _ComplexOp[_NBit1] + __rpow__: _ComplexOp[_NBit1] + +complex64 = complexfloating[_32Bit, _32Bit] +complex128 = complexfloating[_64Bit, _64Bit] + +csingle = complexfloating[_NBitSingle, _NBitSingle] +singlecomplex = complexfloating[_NBitSingle, _NBitSingle] +cdouble = complexfloating[_NBitDouble, _NBitDouble] +complex_ = complexfloating[_NBitDouble, _NBitDouble] +cfloat = complexfloating[_NBitDouble, _NBitDouble] +clongdouble = complexfloating[_NBitLongDouble, _NBitLongDouble] +clongfloat = complexfloating[_NBitLongDouble, _NBitLongDouble] +longcomplex = complexfloating[_NBitLongDouble, _NBitLongDouble] + +class flexible(generic): ... # type: ignore + +# TODO: `item`/`tolist` returns either `bytes` or `tuple` +# depending on whether or not it's used as an opaque bytes sequence +# or a structure +class void(flexible): + @overload + def __init__(self, value: _IntLike_co | bytes, /, dtype : None = ...) -> None: ... + @overload + def __init__(self, value: Any, /, dtype: _DTypeLikeVoid) -> None: ... + @property + def real(self: _ArraySelf) -> _ArraySelf: ... + @property + def imag(self: _ArraySelf) -> _ArraySelf: ... + def setfield( + self, val: ArrayLike, dtype: DTypeLike, offset: int = ... + ) -> None: ... + @overload + def __getitem__(self, key: str | SupportsIndex) -> Any: ... + @overload + def __getitem__(self, key: list[str]) -> void: ... + def __setitem__( + self, + key: str | list[str] | SupportsIndex, + value: ArrayLike, + ) -> None: ... + +class character(flexible): # type: ignore + def __int__(self) -> int: ... + def __float__(self) -> float: ... + +# NOTE: Most `np.bytes_` / `np.str_` methods return their +# builtin `bytes` / `str` counterpart + +class bytes_(character, bytes): + @overload + def __init__(self, value: object = ..., /) -> None: ... + @overload + def __init__( + self, value: str, /, encoding: str = ..., errors: str = ... + ) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> bytes: ... + def tolist(self) -> bytes: ... + +string_ = bytes_ + +class str_(character, str): + @overload + def __init__(self, value: object = ..., /) -> None: ... + @overload + def __init__( + self, value: bytes, /, encoding: str = ..., errors: str = ... + ) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> str: ... + def tolist(self) -> str: ... + +unicode_ = str_ + +# +# Constants +# + +Inf: Final[float] +Infinity: Final[float] +NAN: Final[float] +NINF: Final[float] +NZERO: Final[float] +NaN: Final[float] +PINF: Final[float] +PZERO: Final[float] +e: Final[float] +euler_gamma: Final[float] +inf: Final[float] +infty: Final[float] +nan: Final[float] +pi: Final[float] + +ERR_IGNORE: L[0] +ERR_WARN: L[1] +ERR_RAISE: L[2] +ERR_CALL: L[3] +ERR_PRINT: L[4] +ERR_LOG: L[5] +ERR_DEFAULT: L[521] + +SHIFT_DIVIDEBYZERO: L[0] +SHIFT_OVERFLOW: L[3] +SHIFT_UNDERFLOW: L[6] +SHIFT_INVALID: L[9] + +FPE_DIVIDEBYZERO: L[1] +FPE_OVERFLOW: L[2] +FPE_UNDERFLOW: L[4] +FPE_INVALID: L[8] + +FLOATING_POINT_SUPPORT: L[1] +UFUNC_BUFSIZE_DEFAULT = BUFSIZE + +little_endian: Final[bool] +True_: Final[bool_] +False_: Final[bool_] + +UFUNC_PYVALS_NAME: L["UFUNC_PYVALS"] + +newaxis: None + +# See `numpy._typing._ufunc` for more concrete nin-/nout-specific stubs +@final +class ufunc: + @property + def __name__(self) -> str: ... + @property + def __doc__(self) -> str: ... + __call__: Callable[..., Any] + @property + def nin(self) -> int: ... + @property + def nout(self) -> int: ... + @property + def nargs(self) -> int: ... + @property + def ntypes(self) -> int: ... + @property + def types(self) -> list[str]: ... + # Broad return type because it has to encompass things like + # + # >>> np.logical_and.identity is True + # True + # >>> np.add.identity is 0 + # True + # >>> np.sin.identity is None + # True + # + # and any user-defined ufuncs. + @property + def identity(self) -> Any: ... + # This is None for ufuncs and a string for gufuncs. + @property + def signature(self) -> None | str: ... + # The next four methods will always exist, but they will just + # raise a ValueError ufuncs with that don't accept two input + # arguments and return one output argument. Because of that we + # can't type them very precisely. + reduce: Any + accumulate: Any + reduceat: Any + outer: Any + # Similarly at won't be defined for ufuncs that return multiple + # outputs, so we can't type it very precisely. + at: Any + +# Parameters: `__name__`, `ntypes` and `identity` +absolute: _UFunc_Nin1_Nout1[L['absolute'], L[20], None] +add: _UFunc_Nin2_Nout1[L['add'], L[22], L[0]] +arccos: _UFunc_Nin1_Nout1[L['arccos'], L[8], None] +arccosh: _UFunc_Nin1_Nout1[L['arccosh'], L[8], None] +arcsin: _UFunc_Nin1_Nout1[L['arcsin'], L[8], None] +arcsinh: _UFunc_Nin1_Nout1[L['arcsinh'], L[8], None] +arctan2: _UFunc_Nin2_Nout1[L['arctan2'], L[5], None] +arctan: _UFunc_Nin1_Nout1[L['arctan'], L[8], None] +arctanh: _UFunc_Nin1_Nout1[L['arctanh'], L[8], None] +bitwise_and: _UFunc_Nin2_Nout1[L['bitwise_and'], L[12], L[-1]] +bitwise_not: _UFunc_Nin1_Nout1[L['invert'], L[12], None] +bitwise_or: _UFunc_Nin2_Nout1[L['bitwise_or'], L[12], L[0]] +bitwise_xor: _UFunc_Nin2_Nout1[L['bitwise_xor'], L[12], L[0]] +cbrt: _UFunc_Nin1_Nout1[L['cbrt'], L[5], None] +ceil: _UFunc_Nin1_Nout1[L['ceil'], L[7], None] +conj: _UFunc_Nin1_Nout1[L['conjugate'], L[18], None] +conjugate: _UFunc_Nin1_Nout1[L['conjugate'], L[18], None] +copysign: _UFunc_Nin2_Nout1[L['copysign'], L[4], None] +cos: _UFunc_Nin1_Nout1[L['cos'], L[9], None] +cosh: _UFunc_Nin1_Nout1[L['cosh'], L[8], None] +deg2rad: _UFunc_Nin1_Nout1[L['deg2rad'], L[5], None] +degrees: _UFunc_Nin1_Nout1[L['degrees'], L[5], None] +divide: _UFunc_Nin2_Nout1[L['true_divide'], L[11], None] +divmod: _UFunc_Nin2_Nout2[L['divmod'], L[15], None] +equal: _UFunc_Nin2_Nout1[L['equal'], L[23], None] +exp2: _UFunc_Nin1_Nout1[L['exp2'], L[8], None] +exp: _UFunc_Nin1_Nout1[L['exp'], L[10], None] +expm1: _UFunc_Nin1_Nout1[L['expm1'], L[8], None] +fabs: _UFunc_Nin1_Nout1[L['fabs'], L[5], None] +float_power: _UFunc_Nin2_Nout1[L['float_power'], L[4], None] +floor: _UFunc_Nin1_Nout1[L['floor'], L[7], None] +floor_divide: _UFunc_Nin2_Nout1[L['floor_divide'], L[21], None] +fmax: _UFunc_Nin2_Nout1[L['fmax'], L[21], None] +fmin: _UFunc_Nin2_Nout1[L['fmin'], L[21], None] +fmod: _UFunc_Nin2_Nout1[L['fmod'], L[15], None] +frexp: _UFunc_Nin1_Nout2[L['frexp'], L[4], None] +gcd: _UFunc_Nin2_Nout1[L['gcd'], L[11], L[0]] +greater: _UFunc_Nin2_Nout1[L['greater'], L[23], None] +greater_equal: _UFunc_Nin2_Nout1[L['greater_equal'], L[23], None] +heaviside: _UFunc_Nin2_Nout1[L['heaviside'], L[4], None] +hypot: _UFunc_Nin2_Nout1[L['hypot'], L[5], L[0]] +invert: _UFunc_Nin1_Nout1[L['invert'], L[12], None] +isfinite: _UFunc_Nin1_Nout1[L['isfinite'], L[20], None] +isinf: _UFunc_Nin1_Nout1[L['isinf'], L[20], None] +isnan: _UFunc_Nin1_Nout1[L['isnan'], L[20], None] +isnat: _UFunc_Nin1_Nout1[L['isnat'], L[2], None] +lcm: _UFunc_Nin2_Nout1[L['lcm'], L[11], None] +ldexp: _UFunc_Nin2_Nout1[L['ldexp'], L[8], None] +left_shift: _UFunc_Nin2_Nout1[L['left_shift'], L[11], None] +less: _UFunc_Nin2_Nout1[L['less'], L[23], None] +less_equal: _UFunc_Nin2_Nout1[L['less_equal'], L[23], None] +log10: _UFunc_Nin1_Nout1[L['log10'], L[8], None] +log1p: _UFunc_Nin1_Nout1[L['log1p'], L[8], None] +log2: _UFunc_Nin1_Nout1[L['log2'], L[8], None] +log: _UFunc_Nin1_Nout1[L['log'], L[10], None] +logaddexp2: _UFunc_Nin2_Nout1[L['logaddexp2'], L[4], float] +logaddexp: _UFunc_Nin2_Nout1[L['logaddexp'], L[4], float] +logical_and: _UFunc_Nin2_Nout1[L['logical_and'], L[20], L[True]] +logical_not: _UFunc_Nin1_Nout1[L['logical_not'], L[20], None] +logical_or: _UFunc_Nin2_Nout1[L['logical_or'], L[20], L[False]] +logical_xor: _UFunc_Nin2_Nout1[L['logical_xor'], L[19], L[False]] +matmul: _GUFunc_Nin2_Nout1[L['matmul'], L[19], None] +maximum: _UFunc_Nin2_Nout1[L['maximum'], L[21], None] +minimum: _UFunc_Nin2_Nout1[L['minimum'], L[21], None] +mod: _UFunc_Nin2_Nout1[L['remainder'], L[16], None] +modf: _UFunc_Nin1_Nout2[L['modf'], L[4], None] +multiply: _UFunc_Nin2_Nout1[L['multiply'], L[23], L[1]] +negative: _UFunc_Nin1_Nout1[L['negative'], L[19], None] +nextafter: _UFunc_Nin2_Nout1[L['nextafter'], L[4], None] +not_equal: _UFunc_Nin2_Nout1[L['not_equal'], L[23], None] +positive: _UFunc_Nin1_Nout1[L['positive'], L[19], None] +power: _UFunc_Nin2_Nout1[L['power'], L[18], None] +rad2deg: _UFunc_Nin1_Nout1[L['rad2deg'], L[5], None] +radians: _UFunc_Nin1_Nout1[L['radians'], L[5], None] +reciprocal: _UFunc_Nin1_Nout1[L['reciprocal'], L[18], None] +remainder: _UFunc_Nin2_Nout1[L['remainder'], L[16], None] +right_shift: _UFunc_Nin2_Nout1[L['right_shift'], L[11], None] +rint: _UFunc_Nin1_Nout1[L['rint'], L[10], None] +sign: _UFunc_Nin1_Nout1[L['sign'], L[19], None] +signbit: _UFunc_Nin1_Nout1[L['signbit'], L[4], None] +sin: _UFunc_Nin1_Nout1[L['sin'], L[9], None] +sinh: _UFunc_Nin1_Nout1[L['sinh'], L[8], None] +spacing: _UFunc_Nin1_Nout1[L['spacing'], L[4], None] +sqrt: _UFunc_Nin1_Nout1[L['sqrt'], L[10], None] +square: _UFunc_Nin1_Nout1[L['square'], L[18], None] +subtract: _UFunc_Nin2_Nout1[L['subtract'], L[21], None] +tan: _UFunc_Nin1_Nout1[L['tan'], L[8], None] +tanh: _UFunc_Nin1_Nout1[L['tanh'], L[8], None] +true_divide: _UFunc_Nin2_Nout1[L['true_divide'], L[11], None] +trunc: _UFunc_Nin1_Nout1[L['trunc'], L[7], None] + +abs = absolute + +class _CopyMode(enum.Enum): + ALWAYS: L[True] + IF_NEEDED: L[False] + NEVER: L[2] + +# Warnings +class RankWarning(UserWarning): ... + +_CallType = TypeVar("_CallType", bound=_ErrFunc | _SupportsWrite[str]) + +class errstate(Generic[_CallType], ContextDecorator): + call: _CallType + kwargs: _ErrDictOptional + + # Expand `**kwargs` into explicit keyword-only arguments + def __init__( + self, + *, + call: _CallType = ..., + all: None | _ErrKind = ..., + divide: None | _ErrKind = ..., + over: None | _ErrKind = ..., + under: None | _ErrKind = ..., + invalid: None | _ErrKind = ..., + ) -> None: ... + def __enter__(self) -> None: ... + def __exit__( + self, + exc_type: None | type[BaseException], + exc_value: None | BaseException, + traceback: None | TracebackType, + /, + ) -> None: ... + +@contextmanager +def _no_nep50_warning() -> Generator[None, None, None]: ... +def _get_promotion_state() -> str: ... +def _set_promotion_state(state: str, /) -> None: ... + +class ndenumerate(Generic[_ScalarType]): + iter: flatiter[NDArray[_ScalarType]] + @overload + def __new__( + cls, arr: _FiniteNestedSequence[_SupportsArray[dtype[_ScalarType]]], + ) -> ndenumerate[_ScalarType]: ... + @overload + def __new__(cls, arr: str | _NestedSequence[str]) -> ndenumerate[str_]: ... + @overload + def __new__(cls, arr: bytes | _NestedSequence[bytes]) -> ndenumerate[bytes_]: ... + @overload + def __new__(cls, arr: bool | _NestedSequence[bool]) -> ndenumerate[bool_]: ... + @overload + def __new__(cls, arr: int | _NestedSequence[int]) -> ndenumerate[int_]: ... + @overload + def __new__(cls, arr: float | _NestedSequence[float]) -> ndenumerate[float_]: ... + @overload + def __new__(cls, arr: complex | _NestedSequence[complex]) -> ndenumerate[complex_]: ... + def __next__(self: ndenumerate[_ScalarType]) -> tuple[_Shape, _ScalarType]: ... + def __iter__(self: _T) -> _T: ... + +class ndindex: + @overload + def __init__(self, shape: tuple[SupportsIndex, ...], /) -> None: ... + @overload + def __init__(self, *shape: SupportsIndex) -> None: ... + def __iter__(self: _T) -> _T: ... + def __next__(self) -> _Shape: ... + +class DataSource: + def __init__( + self, + destpath: None | str | os.PathLike[str] = ..., + ) -> None: ... + def __del__(self) -> None: ... + def abspath(self, path: str) -> str: ... + def exists(self, path: str) -> bool: ... + + # Whether the file-object is opened in string or bytes mode (by default) + # depends on the file-extension of `path` + def open( + self, + path: str, + mode: str = ..., + encoding: None | str = ..., + newline: None | str = ..., + ) -> IO[Any]: ... + +# TODO: The type of each `__next__` and `iters` return-type depends +# on the length and dtype of `args`; we can't describe this behavior yet +# as we lack variadics (PEP 646). +@final +class broadcast: + def __new__(cls, *args: ArrayLike) -> broadcast: ... + @property + def index(self) -> int: ... + @property + def iters(self) -> tuple[flatiter[Any], ...]: ... + @property + def nd(self) -> int: ... + @property + def ndim(self) -> int: ... + @property + def numiter(self) -> int: ... + @property + def shape(self) -> _Shape: ... + @property + def size(self) -> int: ... + def __next__(self) -> tuple[Any, ...]: ... + def __iter__(self: _T) -> _T: ... + def reset(self) -> None: ... + +@final +class busdaycalendar: + def __new__( + cls, + weekmask: ArrayLike = ..., + holidays: ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + ) -> busdaycalendar: ... + @property + def weekmask(self) -> NDArray[bool_]: ... + @property + def holidays(self) -> NDArray[datetime64]: ... + +class finfo(Generic[_FloatType]): + dtype: dtype[_FloatType] + bits: int + eps: _FloatType + epsneg: _FloatType + iexp: int + machep: int + max: _FloatType + maxexp: int + min: _FloatType + minexp: int + negep: int + nexp: int + nmant: int + precision: int + resolution: _FloatType + smallest_subnormal: _FloatType + @property + def smallest_normal(self) -> _FloatType: ... + @property + def tiny(self) -> _FloatType: ... + @overload + def __new__( + cls, dtype: inexact[_NBit1] | _DTypeLike[inexact[_NBit1]] + ) -> finfo[floating[_NBit1]]: ... + @overload + def __new__( + cls, dtype: complex | float | type[complex] | type[float] + ) -> finfo[float_]: ... + @overload + def __new__( + cls, dtype: str + ) -> finfo[floating[Any]]: ... + +class iinfo(Generic[_IntType]): + dtype: dtype[_IntType] + kind: str + bits: int + key: str + @property + def min(self) -> int: ... + @property + def max(self) -> int: ... + + @overload + def __new__(cls, dtype: _IntType | _DTypeLike[_IntType]) -> iinfo[_IntType]: ... + @overload + def __new__(cls, dtype: int | type[int]) -> iinfo[int_]: ... + @overload + def __new__(cls, dtype: str) -> iinfo[Any]: ... + +class format_parser: + dtype: dtype[void] + def __init__( + self, + formats: DTypeLike, + names: None | str | Sequence[str], + titles: None | str | Sequence[str], + aligned: bool = ..., + byteorder: None | _ByteOrder = ..., + ) -> None: ... + +class recarray(ndarray[_ShapeType, _DType_co]): + # NOTE: While not strictly mandatory, we're demanding here that arguments + # for the `format_parser`- and `dtype`-based dtype constructors are + # mutually exclusive + @overload + def __new__( + subtype, + shape: _ShapeLike, + dtype: None = ..., + buf: None | _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: None | _ShapeLike = ..., + *, + formats: DTypeLike, + names: None | str | Sequence[str] = ..., + titles: None | str | Sequence[str] = ..., + byteorder: None | _ByteOrder = ..., + aligned: bool = ..., + order: _OrderKACF = ..., + ) -> recarray[Any, dtype[record]]: ... + @overload + def __new__( + subtype, + shape: _ShapeLike, + dtype: DTypeLike, + buf: None | _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: None | _ShapeLike = ..., + formats: None = ..., + names: None = ..., + titles: None = ..., + byteorder: None = ..., + aligned: L[False] = ..., + order: _OrderKACF = ..., + ) -> recarray[Any, dtype[Any]]: ... + def __array_finalize__(self, obj: object) -> None: ... + def __getattribute__(self, attr: str) -> Any: ... + def __setattr__(self, attr: str, val: ArrayLike) -> None: ... + @overload + def __getitem__(self, indx: ( + SupportsIndex + | _ArrayLikeInt_co + | tuple[SupportsIndex | _ArrayLikeInt_co, ...] + )) -> Any: ... + @overload + def __getitem__(self: recarray[Any, dtype[void]], indx: ( + None + | slice + | ellipsis + | SupportsIndex + | _ArrayLikeInt_co + | tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...] + )) -> recarray[Any, _DType_co]: ... + @overload + def __getitem__(self, indx: ( + None + | slice + | ellipsis + | SupportsIndex + | _ArrayLikeInt_co + | tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...] + )) -> ndarray[Any, _DType_co]: ... + @overload + def __getitem__(self, indx: str) -> NDArray[Any]: ... + @overload + def __getitem__(self, indx: list[str]) -> recarray[_ShapeType, dtype[record]]: ... + @overload + def field(self, attr: int | str, val: None = ...) -> Any: ... + @overload + def field(self, attr: int | str, val: ArrayLike) -> None: ... + +class record(void): + def __getattribute__(self, attr: str) -> Any: ... + def __setattr__(self, attr: str, val: ArrayLike) -> None: ... + def pprint(self) -> str: ... + @overload + def __getitem__(self, key: str | SupportsIndex) -> Any: ... + @overload + def __getitem__(self, key: list[str]) -> record: ... + +_NDIterFlagsKind = L[ + "buffered", + "c_index", + "copy_if_overlap", + "common_dtype", + "delay_bufalloc", + "external_loop", + "f_index", + "grow_inner", "growinner", + "multi_index", + "ranged", + "refs_ok", + "reduce_ok", + "zerosize_ok", +] + +_NDIterOpFlagsKind = L[ + "aligned", + "allocate", + "arraymask", + "copy", + "config", + "nbo", + "no_subtype", + "no_broadcast", + "overlap_assume_elementwise", + "readonly", + "readwrite", + "updateifcopy", + "virtual", + "writeonly", + "writemasked" +] + +@final +class nditer: + def __new__( + cls, + op: ArrayLike | Sequence[ArrayLike], + flags: None | Sequence[_NDIterFlagsKind] = ..., + op_flags: None | Sequence[Sequence[_NDIterOpFlagsKind]] = ..., + op_dtypes: DTypeLike | Sequence[DTypeLike] = ..., + order: _OrderKACF = ..., + casting: _CastingKind = ..., + op_axes: None | Sequence[Sequence[SupportsIndex]] = ..., + itershape: None | _ShapeLike = ..., + buffersize: SupportsIndex = ..., + ) -> nditer: ... + def __enter__(self) -> nditer: ... + def __exit__( + self, + exc_type: None | type[BaseException], + exc_value: None | BaseException, + traceback: None | TracebackType, + ) -> None: ... + def __iter__(self) -> nditer: ... + def __next__(self) -> tuple[NDArray[Any], ...]: ... + def __len__(self) -> int: ... + def __copy__(self) -> nditer: ... + @overload + def __getitem__(self, index: SupportsIndex) -> NDArray[Any]: ... + @overload + def __getitem__(self, index: slice) -> tuple[NDArray[Any], ...]: ... + def __setitem__(self, index: slice | SupportsIndex, value: ArrayLike) -> None: ... + def close(self) -> None: ... + def copy(self) -> nditer: ... + def debug_print(self) -> None: ... + def enable_external_loop(self) -> None: ... + def iternext(self) -> bool: ... + def remove_axis(self, i: SupportsIndex, /) -> None: ... + def remove_multi_index(self) -> None: ... + def reset(self) -> None: ... + @property + def dtypes(self) -> tuple[dtype[Any], ...]: ... + @property + def finished(self) -> bool: ... + @property + def has_delayed_bufalloc(self) -> bool: ... + @property + def has_index(self) -> bool: ... + @property + def has_multi_index(self) -> bool: ... + @property + def index(self) -> int: ... + @property + def iterationneedsapi(self) -> bool: ... + @property + def iterindex(self) -> int: ... + @property + def iterrange(self) -> tuple[int, ...]: ... + @property + def itersize(self) -> int: ... + @property + def itviews(self) -> tuple[NDArray[Any], ...]: ... + @property + def multi_index(self) -> tuple[int, ...]: ... + @property + def ndim(self) -> int: ... + @property + def nop(self) -> int: ... + @property + def operands(self) -> tuple[NDArray[Any], ...]: ... + @property + def shape(self) -> tuple[int, ...]: ... + @property + def value(self) -> tuple[NDArray[Any], ...]: ... + +_MemMapModeKind = L[ + "readonly", "r", + "copyonwrite", "c", + "readwrite", "r+", + "write", "w+", +] + +class memmap(ndarray[_ShapeType, _DType_co]): + __array_priority__: ClassVar[float] + filename: str | None + offset: int + mode: str + @overload + def __new__( + subtype, + filename: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _MemMapIOProtocol, + dtype: type[uint8] = ..., + mode: _MemMapModeKind = ..., + offset: int = ..., + shape: None | int | tuple[int, ...] = ..., + order: _OrderKACF = ..., + ) -> memmap[Any, dtype[uint8]]: ... + @overload + def __new__( + subtype, + filename: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _MemMapIOProtocol, + dtype: _DTypeLike[_ScalarType], + mode: _MemMapModeKind = ..., + offset: int = ..., + shape: None | int | tuple[int, ...] = ..., + order: _OrderKACF = ..., + ) -> memmap[Any, dtype[_ScalarType]]: ... + @overload + def __new__( + subtype, + filename: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _MemMapIOProtocol, + dtype: DTypeLike, + mode: _MemMapModeKind = ..., + offset: int = ..., + shape: None | int | tuple[int, ...] = ..., + order: _OrderKACF = ..., + ) -> memmap[Any, dtype[Any]]: ... + def __array_finalize__(self, obj: object) -> None: ... + def __array_wrap__( + self, + array: memmap[_ShapeType, _DType_co], + context: None | tuple[ufunc, tuple[Any, ...], int] = ..., + ) -> Any: ... + def flush(self) -> None: ... + +# TODO: Add a mypy plugin for managing functions whose output type is dependent +# on the literal value of some sort of signature (e.g. `einsum` and `vectorize`) +class vectorize: + pyfunc: Callable[..., Any] + cache: bool + signature: None | str + otypes: None | str + excluded: set[int | str] + __doc__: None | str + def __init__( + self, + pyfunc: Callable[..., Any], + otypes: None | str | Iterable[DTypeLike] = ..., + doc: None | str = ..., + excluded: None | Iterable[int | str] = ..., + cache: bool = ..., + signature: None | str = ..., + ) -> None: ... + def __call__(self, *args: Any, **kwargs: Any) -> Any: ... + +class poly1d: + @property + def variable(self) -> str: ... + @property + def order(self) -> int: ... + @property + def o(self) -> int: ... + @property + def roots(self) -> NDArray[Any]: ... + @property + def r(self) -> NDArray[Any]: ... + + @property + def coeffs(self) -> NDArray[Any]: ... + @coeffs.setter + def coeffs(self, value: NDArray[Any]) -> None: ... + + @property + def c(self) -> NDArray[Any]: ... + @c.setter + def c(self, value: NDArray[Any]) -> None: ... + + @property + def coef(self) -> NDArray[Any]: ... + @coef.setter + def coef(self, value: NDArray[Any]) -> None: ... + + @property + def coefficients(self) -> NDArray[Any]: ... + @coefficients.setter + def coefficients(self, value: NDArray[Any]) -> None: ... + + __hash__: ClassVar[None] # type: ignore + + @overload + def __array__(self, t: None = ...) -> NDArray[Any]: ... + @overload + def __array__(self, t: _DType) -> ndarray[Any, _DType]: ... + + @overload + def __call__(self, val: _ScalarLike_co) -> Any: ... + @overload + def __call__(self, val: poly1d) -> poly1d: ... + @overload + def __call__(self, val: ArrayLike) -> NDArray[Any]: ... + + def __init__( + self, + c_or_r: ArrayLike, + r: bool = ..., + variable: None | str = ..., + ) -> None: ... + def __len__(self) -> int: ... + def __neg__(self) -> poly1d: ... + def __pos__(self) -> poly1d: ... + def __mul__(self, other: ArrayLike) -> poly1d: ... + def __rmul__(self, other: ArrayLike) -> poly1d: ... + def __add__(self, other: ArrayLike) -> poly1d: ... + def __radd__(self, other: ArrayLike) -> poly1d: ... + def __pow__(self, val: _FloatLike_co) -> poly1d: ... # Integral floats are accepted + def __sub__(self, other: ArrayLike) -> poly1d: ... + def __rsub__(self, other: ArrayLike) -> poly1d: ... + def __div__(self, other: ArrayLike) -> poly1d: ... + def __truediv__(self, other: ArrayLike) -> poly1d: ... + def __rdiv__(self, other: ArrayLike) -> poly1d: ... + def __rtruediv__(self, other: ArrayLike) -> poly1d: ... + def __getitem__(self, val: int) -> Any: ... + def __setitem__(self, key: int, val: Any) -> None: ... + def __iter__(self) -> Iterator[Any]: ... + def deriv(self, m: SupportsInt | SupportsIndex = ...) -> poly1d: ... + def integ( + self, + m: SupportsInt | SupportsIndex = ..., + k: None | _ArrayLikeComplex_co | _ArrayLikeObject_co = ..., + ) -> poly1d: ... + +class matrix(ndarray[_ShapeType, _DType_co]): + __array_priority__: ClassVar[float] + def __new__( + subtype, + data: ArrayLike, + dtype: DTypeLike = ..., + copy: bool = ..., + ) -> matrix[Any, Any]: ... + def __array_finalize__(self, obj: object) -> None: ... + + @overload + def __getitem__(self, key: ( + SupportsIndex + | _ArrayLikeInt_co + | tuple[SupportsIndex | _ArrayLikeInt_co, ...] + )) -> Any: ... + @overload + def __getitem__(self, key: ( + None + | slice + | ellipsis + | SupportsIndex + | _ArrayLikeInt_co + | tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...] + )) -> matrix[Any, _DType_co]: ... + @overload + def __getitem__(self: NDArray[void], key: str) -> matrix[Any, dtype[Any]]: ... + @overload + def __getitem__(self: NDArray[void], key: list[str]) -> matrix[_ShapeType, dtype[void]]: ... + + def __mul__(self, other: ArrayLike) -> matrix[Any, Any]: ... + def __rmul__(self, other: ArrayLike) -> matrix[Any, Any]: ... + def __imul__(self, other: ArrayLike) -> matrix[_ShapeType, _DType_co]: ... + def __pow__(self, other: ArrayLike) -> matrix[Any, Any]: ... + def __ipow__(self, other: ArrayLike) -> matrix[_ShapeType, _DType_co]: ... + + @overload + def sum(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ... + @overload + def sum(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ...) -> matrix[Any, Any]: ... + @overload + def sum(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def mean(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ... + @overload + def mean(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ...) -> matrix[Any, Any]: ... + @overload + def mean(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def std(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> Any: ... + @overload + def std(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> matrix[Any, Any]: ... + @overload + def std(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ..., ddof: float = ...) -> _NdArraySubClass: ... + + @overload + def var(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> Any: ... + @overload + def var(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> matrix[Any, Any]: ... + @overload + def var(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ..., ddof: float = ...) -> _NdArraySubClass: ... + + @overload + def prod(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ... + @overload + def prod(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ...) -> matrix[Any, Any]: ... + @overload + def prod(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def any(self, axis: None = ..., out: None = ...) -> bool_: ... + @overload + def any(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[bool_]]: ... + @overload + def any(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def all(self, axis: None = ..., out: None = ...) -> bool_: ... + @overload + def all(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[bool_]]: ... + @overload + def all(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def max(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> _ScalarType: ... + @overload + def max(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, _DType_co]: ... + @overload + def max(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def min(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> _ScalarType: ... + @overload + def min(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, _DType_co]: ... + @overload + def min(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def argmax(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> intp: ... + @overload + def argmax(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[intp]]: ... + @overload + def argmax(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def argmin(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> intp: ... + @overload + def argmin(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[intp]]: ... + @overload + def argmin(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def ptp(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> _ScalarType: ... + @overload + def ptp(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, _DType_co]: ... + @overload + def ptp(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + def squeeze(self, axis: None | _ShapeLike = ...) -> matrix[Any, _DType_co]: ... + def tolist(self: matrix[Any, dtype[_SupportsItem[_T]]]) -> list[list[_T]]: ... # type: ignore[typevar] + def ravel(self, order: _OrderKACF = ...) -> matrix[Any, _DType_co]: ... + def flatten(self, order: _OrderKACF = ...) -> matrix[Any, _DType_co]: ... + + @property + def T(self) -> matrix[Any, _DType_co]: ... + @property + def I(self) -> matrix[Any, Any]: ... + @property + def A(self) -> ndarray[_ShapeType, _DType_co]: ... + @property + def A1(self) -> ndarray[Any, _DType_co]: ... + @property + def H(self) -> matrix[Any, _DType_co]: ... + def getT(self) -> matrix[Any, _DType_co]: ... + def getI(self) -> matrix[Any, Any]: ... + def getA(self) -> ndarray[_ShapeType, _DType_co]: ... + def getA1(self) -> ndarray[Any, _DType_co]: ... + def getH(self) -> matrix[Any, _DType_co]: ... + +_CharType = TypeVar("_CharType", str_, bytes_) +_CharDType = TypeVar("_CharDType", dtype[str_], dtype[bytes_]) +_CharArray = chararray[Any, dtype[_CharType]] + +class chararray(ndarray[_ShapeType, _CharDType]): + @overload + def __new__( + subtype, + shape: _ShapeLike, + itemsize: SupportsIndex | SupportsInt = ..., + unicode: L[False] = ..., + buffer: _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: _ShapeLike = ..., + order: _OrderKACF = ..., + ) -> chararray[Any, dtype[bytes_]]: ... + @overload + def __new__( + subtype, + shape: _ShapeLike, + itemsize: SupportsIndex | SupportsInt = ..., + unicode: L[True] = ..., + buffer: _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: _ShapeLike = ..., + order: _OrderKACF = ..., + ) -> chararray[Any, dtype[str_]]: ... + + def __array_finalize__(self, obj: object) -> None: ... + def __mul__(self, other: _ArrayLikeInt_co) -> chararray[Any, _CharDType]: ... + def __rmul__(self, other: _ArrayLikeInt_co) -> chararray[Any, _CharDType]: ... + def __mod__(self, i: Any) -> chararray[Any, _CharDType]: ... + + @overload + def __eq__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __eq__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __ne__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __ne__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __ge__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __ge__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __le__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __le__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __gt__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __gt__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __lt__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __lt__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __add__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def __add__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def __radd__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def __radd__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def center( + self: _CharArray[str_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def center( + self: _CharArray[bytes_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def count( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def count( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + def decode( + self: _CharArray[bytes_], + encoding: None | str = ..., + errors: None | str = ..., + ) -> _CharArray[str_]: ... + + def encode( + self: _CharArray[str_], + encoding: None | str = ..., + errors: None | str = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def endswith( + self: _CharArray[str_], + suffix: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[bool_]: ... + @overload + def endswith( + self: _CharArray[bytes_], + suffix: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[bool_]: ... + + def expandtabs( + self, + tabsize: _ArrayLikeInt_co = ..., + ) -> chararray[Any, _CharDType]: ... + + @overload + def find( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def find( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + @overload + def index( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def index( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + @overload + def join( + self: _CharArray[str_], + seq: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def join( + self: _CharArray[bytes_], + seq: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def ljust( + self: _CharArray[str_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def ljust( + self: _CharArray[bytes_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def lstrip( + self: _CharArray[str_], + chars: None | _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def lstrip( + self: _CharArray[bytes_], + chars: None | _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def partition( + self: _CharArray[str_], + sep: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def partition( + self: _CharArray[bytes_], + sep: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def replace( + self: _CharArray[str_], + old: _ArrayLikeStr_co, + new: _ArrayLikeStr_co, + count: None | _ArrayLikeInt_co = ..., + ) -> _CharArray[str_]: ... + @overload + def replace( + self: _CharArray[bytes_], + old: _ArrayLikeBytes_co, + new: _ArrayLikeBytes_co, + count: None | _ArrayLikeInt_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def rfind( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def rfind( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + @overload + def rindex( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def rindex( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + @overload + def rjust( + self: _CharArray[str_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def rjust( + self: _CharArray[bytes_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def rpartition( + self: _CharArray[str_], + sep: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def rpartition( + self: _CharArray[bytes_], + sep: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def rsplit( + self: _CharArray[str_], + sep: None | _ArrayLikeStr_co = ..., + maxsplit: None | _ArrayLikeInt_co = ..., + ) -> NDArray[object_]: ... + @overload + def rsplit( + self: _CharArray[bytes_], + sep: None | _ArrayLikeBytes_co = ..., + maxsplit: None | _ArrayLikeInt_co = ..., + ) -> NDArray[object_]: ... + + @overload + def rstrip( + self: _CharArray[str_], + chars: None | _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def rstrip( + self: _CharArray[bytes_], + chars: None | _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def split( + self: _CharArray[str_], + sep: None | _ArrayLikeStr_co = ..., + maxsplit: None | _ArrayLikeInt_co = ..., + ) -> NDArray[object_]: ... + @overload + def split( + self: _CharArray[bytes_], + sep: None | _ArrayLikeBytes_co = ..., + maxsplit: None | _ArrayLikeInt_co = ..., + ) -> NDArray[object_]: ... + + def splitlines(self, keepends: None | _ArrayLikeBool_co = ...) -> NDArray[object_]: ... + + @overload + def startswith( + self: _CharArray[str_], + prefix: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[bool_]: ... + @overload + def startswith( + self: _CharArray[bytes_], + prefix: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[bool_]: ... + + @overload + def strip( + self: _CharArray[str_], + chars: None | _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def strip( + self: _CharArray[bytes_], + chars: None | _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def translate( + self: _CharArray[str_], + table: _ArrayLikeStr_co, + deletechars: None | _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def translate( + self: _CharArray[bytes_], + table: _ArrayLikeBytes_co, + deletechars: None | _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + def zfill(self, width: _ArrayLikeInt_co) -> chararray[Any, _CharDType]: ... + def capitalize(self) -> chararray[_ShapeType, _CharDType]: ... + def title(self) -> chararray[_ShapeType, _CharDType]: ... + def swapcase(self) -> chararray[_ShapeType, _CharDType]: ... + def lower(self) -> chararray[_ShapeType, _CharDType]: ... + def upper(self) -> chararray[_ShapeType, _CharDType]: ... + def isalnum(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isalpha(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isdigit(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def islower(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isspace(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def istitle(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isupper(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isnumeric(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isdecimal(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + +# NOTE: Deprecated +# class MachAr: ... + +class _SupportsDLPack(Protocol[_T_contra]): + def __dlpack__(self, *, stream: None | _T_contra = ...) -> _PyCapsule: ... + +def from_dlpack(obj: _SupportsDLPack[None], /) -> NDArray[Any]: ... diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/_distributor_init.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/_distributor_init.py new file mode 100644 index 0000000000000000000000000000000000000000..25b0eed79fcabe6d6ad5a7b2bf45e5371f37d4a0 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/_distributor_init.py @@ -0,0 +1,15 @@ +""" Distributor init file + +Distributors: you can add custom code here to support particular distributions +of numpy. + +For example, this is a good place to put any BLAS/LAPACK initialization code. + +The numpy standard source distribution will not put code in this file, so you +can safely replace this file with your own version. +""" + +try: + from . import _distributor_init_local +except ImportError: + pass diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/_globals.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/_globals.py new file mode 100644 index 0000000000000000000000000000000000000000..416a20f5e11b14b1da34e2bfb45c7961edc9097c --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/_globals.py @@ -0,0 +1,95 @@ +""" +Module defining global singleton classes. + +This module raises a RuntimeError if an attempt to reload it is made. In that +way the identities of the classes defined here are fixed and will remain so +even if numpy itself is reloaded. In particular, a function like the following +will still work correctly after numpy is reloaded:: + + def foo(arg=np._NoValue): + if arg is np._NoValue: + ... + +That was not the case when the singleton classes were defined in the numpy +``__init__.py`` file. See gh-7844 for a discussion of the reload problem that +motivated this module. + +""" +import enum + +from ._utils import set_module as _set_module + +__all__ = ['_NoValue', '_CopyMode'] + + +# Disallow reloading this module so as to preserve the identities of the +# classes defined here. +if '_is_loaded' in globals(): + raise RuntimeError('Reloading numpy._globals is not allowed') +_is_loaded = True + + +class _NoValueType: + """Special keyword value. + + The instance of this class may be used as the default value assigned to a + keyword if no other obvious default (e.g., `None`) is suitable, + + Common reasons for using this keyword are: + + - A new keyword is added to a function, and that function forwards its + inputs to another function or method which can be defined outside of + NumPy. For example, ``np.std(x)`` calls ``x.std``, so when a ``keepdims`` + keyword was added that could only be forwarded if the user explicitly + specified ``keepdims``; downstream array libraries may not have added + the same keyword, so adding ``x.std(..., keepdims=keepdims)`` + unconditionally could have broken previously working code. + - A keyword is being deprecated, and a deprecation warning must only be + emitted when the keyword is used. + + """ + __instance = None + def __new__(cls): + # ensure that only one instance exists + if not cls.__instance: + cls.__instance = super().__new__(cls) + return cls.__instance + + def __repr__(self): + return "" + + +_NoValue = _NoValueType() + + +@_set_module("numpy") +class _CopyMode(enum.Enum): + """ + An enumeration for the copy modes supported + by numpy.copy() and numpy.array(). The following three modes are supported, + + - ALWAYS: This means that a deep copy of the input + array will always be taken. + - IF_NEEDED: This means that a deep copy of the input + array will be taken only if necessary. + - NEVER: This means that the deep copy will never be taken. + If a copy cannot be avoided then a `ValueError` will be + raised. + + Note that the buffer-protocol could in theory do copies. NumPy currently + assumes an object exporting the buffer protocol will never do this. + """ + + ALWAYS = True + IF_NEEDED = False + NEVER = 2 + + def __bool__(self): + # For backwards compatibility + if self == _CopyMode.ALWAYS: + return True + + if self == _CopyMode.IF_NEEDED: + return False + + raise ValueError(f"{self} is neither True nor False.") diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/_pytesttester.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/_pytesttester.py new file mode 100644 index 0000000000000000000000000000000000000000..1c38291ae3319a08bb665fe5c86dfa13e1655a4c --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/_pytesttester.py @@ -0,0 +1,207 @@ +""" +Pytest test running. + +This module implements the ``test()`` function for NumPy modules. The usual +boiler plate for doing that is to put the following in the module +``__init__.py`` file:: + + from numpy._pytesttester import PytestTester + test = PytestTester(__name__) + del PytestTester + + +Warnings filtering and other runtime settings should be dealt with in the +``pytest.ini`` file in the numpy repo root. The behavior of the test depends on +whether or not that file is found as follows: + +* ``pytest.ini`` is present (develop mode) + All warnings except those explicitly filtered out are raised as error. +* ``pytest.ini`` is absent (release mode) + DeprecationWarnings and PendingDeprecationWarnings are ignored, other + warnings are passed through. + +In practice, tests run from the numpy repo are run in develop mode. That +includes the standard ``python runtests.py`` invocation. + +This module is imported by every numpy subpackage, so lies at the top level to +simplify circular import issues. For the same reason, it contains no numpy +imports at module scope, instead importing numpy within function calls. +""" +import sys +import os + +__all__ = ['PytestTester'] + + +def _show_numpy_info(): + import numpy as np + + print("NumPy version %s" % np.__version__) + relaxed_strides = np.ones((10, 1), order="C").flags.f_contiguous + print("NumPy relaxed strides checking option:", relaxed_strides) + info = np.lib.utils._opt_info() + print("NumPy CPU features: ", (info if info else 'nothing enabled')) + + +class PytestTester: + """ + Pytest test runner. + + A test function is typically added to a package's __init__.py like so:: + + from numpy._pytesttester import PytestTester + test = PytestTester(__name__).test + del PytestTester + + Calling this test function finds and runs all tests associated with the + module and all its sub-modules. + + Attributes + ---------- + module_name : str + Full path to the package to test. + + Parameters + ---------- + module_name : module name + The name of the module to test. + + Notes + ----- + Unlike the previous ``nose``-based implementation, this class is not + publicly exposed as it performs some ``numpy``-specific warning + suppression. + + """ + def __init__(self, module_name): + self.module_name = module_name + + def __call__(self, label='fast', verbose=1, extra_argv=None, + doctests=False, coverage=False, durations=-1, tests=None): + """ + Run tests for module using pytest. + + Parameters + ---------- + label : {'fast', 'full'}, optional + Identifies the tests to run. When set to 'fast', tests decorated + with `pytest.mark.slow` are skipped, when 'full', the slow marker + is ignored. + verbose : int, optional + Verbosity value for test outputs, in the range 1-3. Default is 1. + extra_argv : list, optional + List with any extra arguments to pass to pytests. + doctests : bool, optional + .. note:: Not supported + coverage : bool, optional + If True, report coverage of NumPy code. Default is False. + Requires installation of (pip) pytest-cov. + durations : int, optional + If < 0, do nothing, If 0, report time of all tests, if > 0, + report the time of the slowest `timer` tests. Default is -1. + tests : test or list of tests + Tests to be executed with pytest '--pyargs' + + Returns + ------- + result : bool + Return True on success, false otherwise. + + Notes + ----- + Each NumPy module exposes `test` in its namespace to run all tests for + it. For example, to run all tests for numpy.lib: + + >>> np.lib.test() #doctest: +SKIP + + Examples + -------- + >>> result = np.lib.test() #doctest: +SKIP + ... + 1023 passed, 2 skipped, 6 deselected, 1 xfailed in 10.39 seconds + >>> result + True + + """ + import pytest + import warnings + + module = sys.modules[self.module_name] + module_path = os.path.abspath(module.__path__[0]) + + # setup the pytest arguments + pytest_args = ["-l"] + + # offset verbosity. The "-q" cancels a "-v". + pytest_args += ["-q"] + + if sys.version_info < (3, 12): + with warnings.catch_warnings(): + warnings.simplefilter("always") + # Filter out distutils cpu warnings (could be localized to + # distutils tests). ASV has problems with top level import, + # so fetch module for suppression here. + from numpy.distutils import cpuinfo + + with warnings.catch_warnings(record=True): + # Ignore the warning from importing the array_api submodule. This + # warning is done on import, so it would break pytest collection, + # but importing it early here prevents the warning from being + # issued when it imported again. + import numpy.array_api + + # Filter out annoying import messages. Want these in both develop and + # release mode. + pytest_args += [ + "-W ignore:Not importing directory", + "-W ignore:numpy.dtype size changed", + "-W ignore:numpy.ufunc size changed", + "-W ignore::UserWarning:cpuinfo", + ] + + # When testing matrices, ignore their PendingDeprecationWarnings + pytest_args += [ + "-W ignore:the matrix subclass is not", + "-W ignore:Importing from numpy.matlib is", + ] + + if doctests: + pytest_args += ["--doctest-modules"] + + if extra_argv: + pytest_args += list(extra_argv) + + if verbose > 1: + pytest_args += ["-" + "v"*(verbose - 1)] + + if coverage: + pytest_args += ["--cov=" + module_path] + + if label == "fast": + # not importing at the top level to avoid circular import of module + from numpy.testing import IS_PYPY + if IS_PYPY: + pytest_args += ["-m", "not slow and not slow_pypy"] + else: + pytest_args += ["-m", "not slow"] + + elif label != "full": + pytest_args += ["-m", label] + + if durations >= 0: + pytest_args += ["--durations=%s" % durations] + + if tests is None: + tests = [self.module_name] + + pytest_args += ["--pyargs"] + list(tests) + + # run tests. + _show_numpy_info() + + try: + code = pytest.main(pytest_args) + except SystemExit as exc: + code = exc.code + + return code == 0 diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/ctypeslib.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/ctypeslib.py new file mode 100644 index 0000000000000000000000000000000000000000..d9f64fd9e716830ff33d4d787a0492c65d517603 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/ctypeslib.py @@ -0,0 +1,545 @@ +""" +============================ +``ctypes`` Utility Functions +============================ + +See Also +-------- +load_library : Load a C library. +ndpointer : Array restype/argtype with verification. +as_ctypes : Create a ctypes array from an ndarray. +as_array : Create an ndarray from a ctypes array. + +References +---------- +.. [1] "SciPy Cookbook: ctypes", https://scipy-cookbook.readthedocs.io/items/Ctypes.html + +Examples +-------- +Load the C library: + +>>> _lib = np.ctypeslib.load_library('libmystuff', '.') #doctest: +SKIP + +Our result type, an ndarray that must be of type double, be 1-dimensional +and is C-contiguous in memory: + +>>> array_1d_double = np.ctypeslib.ndpointer( +... dtype=np.double, +... ndim=1, flags='CONTIGUOUS') #doctest: +SKIP + +Our C-function typically takes an array and updates its values +in-place. For example:: + + void foo_func(double* x, int length) + { + int i; + for (i = 0; i < length; i++) { + x[i] = i*i; + } + } + +We wrap it using: + +>>> _lib.foo_func.restype = None #doctest: +SKIP +>>> _lib.foo_func.argtypes = [array_1d_double, c_int] #doctest: +SKIP + +Then, we're ready to call ``foo_func``: + +>>> out = np.empty(15, dtype=np.double) +>>> _lib.foo_func(out, len(out)) #doctest: +SKIP + +""" +__all__ = ['load_library', 'ndpointer', 'c_intp', 'as_ctypes', 'as_array', + 'as_ctypes_type'] + +import os +from numpy import ( + integer, ndarray, dtype as _dtype, asarray, frombuffer +) +from numpy.core.multiarray import _flagdict, flagsobj + +try: + import ctypes +except ImportError: + ctypes = None + +if ctypes is None: + def _dummy(*args, **kwds): + """ + Dummy object that raises an ImportError if ctypes is not available. + + Raises + ------ + ImportError + If ctypes is not available. + + """ + raise ImportError("ctypes is not available.") + load_library = _dummy + as_ctypes = _dummy + as_array = _dummy + from numpy import intp as c_intp + _ndptr_base = object +else: + import numpy.core._internal as nic + c_intp = nic._getintp_ctype() + del nic + _ndptr_base = ctypes.c_void_p + + # Adapted from Albert Strasheim + def load_library(libname, loader_path): + """ + It is possible to load a library using + + >>> lib = ctypes.cdll[] # doctest: +SKIP + + But there are cross-platform considerations, such as library file extensions, + plus the fact Windows will just load the first library it finds with that name. + NumPy supplies the load_library function as a convenience. + + .. versionchanged:: 1.20.0 + Allow libname and loader_path to take any + :term:`python:path-like object`. + + Parameters + ---------- + libname : path-like + Name of the library, which can have 'lib' as a prefix, + but without an extension. + loader_path : path-like + Where the library can be found. + + Returns + ------- + ctypes.cdll[libpath] : library object + A ctypes library object + + Raises + ------ + OSError + If there is no library with the expected extension, or the + library is defective and cannot be loaded. + """ + # Convert path-like objects into strings + libname = os.fsdecode(libname) + loader_path = os.fsdecode(loader_path) + + ext = os.path.splitext(libname)[1] + if not ext: + import sys + import sysconfig + # Try to load library with platform-specific name, otherwise + # default to libname.[so|dll|dylib]. Sometimes, these files are + # built erroneously on non-linux platforms. + base_ext = ".so" + if sys.platform.startswith("darwin"): + base_ext = ".dylib" + elif sys.platform.startswith("win"): + base_ext = ".dll" + libname_ext = [libname + base_ext] + so_ext = sysconfig.get_config_var("EXT_SUFFIX") + if not so_ext == base_ext: + libname_ext.insert(0, libname + so_ext) + else: + libname_ext = [libname] + + loader_path = os.path.abspath(loader_path) + if not os.path.isdir(loader_path): + libdir = os.path.dirname(loader_path) + else: + libdir = loader_path + + for ln in libname_ext: + libpath = os.path.join(libdir, ln) + if os.path.exists(libpath): + try: + return ctypes.cdll[libpath] + except OSError: + ## defective lib file + raise + ## if no successful return in the libname_ext loop: + raise OSError("no file with expected extension") + + +def _num_fromflags(flaglist): + num = 0 + for val in flaglist: + num += _flagdict[val] + return num + +_flagnames = ['C_CONTIGUOUS', 'F_CONTIGUOUS', 'ALIGNED', 'WRITEABLE', + 'OWNDATA', 'WRITEBACKIFCOPY'] +def _flags_fromnum(num): + res = [] + for key in _flagnames: + value = _flagdict[key] + if (num & value): + res.append(key) + return res + + +class _ndptr(_ndptr_base): + @classmethod + def from_param(cls, obj): + if not isinstance(obj, ndarray): + raise TypeError("argument must be an ndarray") + if cls._dtype_ is not None \ + and obj.dtype != cls._dtype_: + raise TypeError("array must have data type %s" % cls._dtype_) + if cls._ndim_ is not None \ + and obj.ndim != cls._ndim_: + raise TypeError("array must have %d dimension(s)" % cls._ndim_) + if cls._shape_ is not None \ + and obj.shape != cls._shape_: + raise TypeError("array must have shape %s" % str(cls._shape_)) + if cls._flags_ is not None \ + and ((obj.flags.num & cls._flags_) != cls._flags_): + raise TypeError("array must have flags %s" % + _flags_fromnum(cls._flags_)) + return obj.ctypes + + +class _concrete_ndptr(_ndptr): + """ + Like _ndptr, but with `_shape_` and `_dtype_` specified. + + Notably, this means the pointer has enough information to reconstruct + the array, which is not generally true. + """ + def _check_retval_(self): + """ + This method is called when this class is used as the .restype + attribute for a shared-library function, to automatically wrap the + pointer into an array. + """ + return self.contents + + @property + def contents(self): + """ + Get an ndarray viewing the data pointed to by this pointer. + + This mirrors the `contents` attribute of a normal ctypes pointer + """ + full_dtype = _dtype((self._dtype_, self._shape_)) + full_ctype = ctypes.c_char * full_dtype.itemsize + buffer = ctypes.cast(self, ctypes.POINTER(full_ctype)).contents + return frombuffer(buffer, dtype=full_dtype).squeeze(axis=0) + + +# Factory for an array-checking class with from_param defined for +# use with ctypes argtypes mechanism +_pointer_type_cache = {} +def ndpointer(dtype=None, ndim=None, shape=None, flags=None): + """ + Array-checking restype/argtypes. + + An ndpointer instance is used to describe an ndarray in restypes + and argtypes specifications. This approach is more flexible than + using, for example, ``POINTER(c_double)``, since several restrictions + can be specified, which are verified upon calling the ctypes function. + These include data type, number of dimensions, shape and flags. If a + given array does not satisfy the specified restrictions, + a ``TypeError`` is raised. + + Parameters + ---------- + dtype : data-type, optional + Array data-type. + ndim : int, optional + Number of array dimensions. + shape : tuple of ints, optional + Array shape. + flags : str or tuple of str + Array flags; may be one or more of: + + - C_CONTIGUOUS / C / CONTIGUOUS + - F_CONTIGUOUS / F / FORTRAN + - OWNDATA / O + - WRITEABLE / W + - ALIGNED / A + - WRITEBACKIFCOPY / X + + Returns + ------- + klass : ndpointer type object + A type object, which is an ``_ndtpr`` instance containing + dtype, ndim, shape and flags information. + + Raises + ------ + TypeError + If a given array does not satisfy the specified restrictions. + + Examples + -------- + >>> clib.somefunc.argtypes = [np.ctypeslib.ndpointer(dtype=np.float64, + ... ndim=1, + ... flags='C_CONTIGUOUS')] + ... #doctest: +SKIP + >>> clib.somefunc(np.array([1, 2, 3], dtype=np.float64)) + ... #doctest: +SKIP + + """ + + # normalize dtype to an Optional[dtype] + if dtype is not None: + dtype = _dtype(dtype) + + # normalize flags to an Optional[int] + num = None + if flags is not None: + if isinstance(flags, str): + flags = flags.split(',') + elif isinstance(flags, (int, integer)): + num = flags + flags = _flags_fromnum(num) + elif isinstance(flags, flagsobj): + num = flags.num + flags = _flags_fromnum(num) + if num is None: + try: + flags = [x.strip().upper() for x in flags] + except Exception as e: + raise TypeError("invalid flags specification") from e + num = _num_fromflags(flags) + + # normalize shape to an Optional[tuple] + if shape is not None: + try: + shape = tuple(shape) + except TypeError: + # single integer -> 1-tuple + shape = (shape,) + + cache_key = (dtype, ndim, shape, num) + + try: + return _pointer_type_cache[cache_key] + except KeyError: + pass + + # produce a name for the new type + if dtype is None: + name = 'any' + elif dtype.names is not None: + name = str(id(dtype)) + else: + name = dtype.str + if ndim is not None: + name += "_%dd" % ndim + if shape is not None: + name += "_"+"x".join(str(x) for x in shape) + if flags is not None: + name += "_"+"_".join(flags) + + if dtype is not None and shape is not None: + base = _concrete_ndptr + else: + base = _ndptr + + klass = type("ndpointer_%s"%name, (base,), + {"_dtype_": dtype, + "_shape_" : shape, + "_ndim_" : ndim, + "_flags_" : num}) + _pointer_type_cache[cache_key] = klass + return klass + + +if ctypes is not None: + def _ctype_ndarray(element_type, shape): + """ Create an ndarray of the given element type and shape """ + for dim in shape[::-1]: + element_type = dim * element_type + # prevent the type name include np.ctypeslib + element_type.__module__ = None + return element_type + + + def _get_scalar_type_map(): + """ + Return a dictionary mapping native endian scalar dtype to ctypes types + """ + ct = ctypes + simple_types = [ + ct.c_byte, ct.c_short, ct.c_int, ct.c_long, ct.c_longlong, + ct.c_ubyte, ct.c_ushort, ct.c_uint, ct.c_ulong, ct.c_ulonglong, + ct.c_float, ct.c_double, + ct.c_bool, + ] + return {_dtype(ctype): ctype for ctype in simple_types} + + + _scalar_type_map = _get_scalar_type_map() + + + def _ctype_from_dtype_scalar(dtype): + # swapping twice ensure that `=` is promoted to <, >, or | + dtype_with_endian = dtype.newbyteorder('S').newbyteorder('S') + dtype_native = dtype.newbyteorder('=') + try: + ctype = _scalar_type_map[dtype_native] + except KeyError as e: + raise NotImplementedError( + "Converting {!r} to a ctypes type".format(dtype) + ) from None + + if dtype_with_endian.byteorder == '>': + ctype = ctype.__ctype_be__ + elif dtype_with_endian.byteorder == '<': + ctype = ctype.__ctype_le__ + + return ctype + + + def _ctype_from_dtype_subarray(dtype): + element_dtype, shape = dtype.subdtype + ctype = _ctype_from_dtype(element_dtype) + return _ctype_ndarray(ctype, shape) + + + def _ctype_from_dtype_structured(dtype): + # extract offsets of each field + field_data = [] + for name in dtype.names: + field_dtype, offset = dtype.fields[name][:2] + field_data.append((offset, name, _ctype_from_dtype(field_dtype))) + + # ctypes doesn't care about field order + field_data = sorted(field_data, key=lambda f: f[0]) + + if len(field_data) > 1 and all(offset == 0 for offset, name, ctype in field_data): + # union, if multiple fields all at address 0 + size = 0 + _fields_ = [] + for offset, name, ctype in field_data: + _fields_.append((name, ctype)) + size = max(size, ctypes.sizeof(ctype)) + + # pad to the right size + if dtype.itemsize != size: + _fields_.append(('', ctypes.c_char * dtype.itemsize)) + + # we inserted manual padding, so always `_pack_` + return type('union', (ctypes.Union,), dict( + _fields_=_fields_, + _pack_=1, + __module__=None, + )) + else: + last_offset = 0 + _fields_ = [] + for offset, name, ctype in field_data: + padding = offset - last_offset + if padding < 0: + raise NotImplementedError("Overlapping fields") + if padding > 0: + _fields_.append(('', ctypes.c_char * padding)) + + _fields_.append((name, ctype)) + last_offset = offset + ctypes.sizeof(ctype) + + + padding = dtype.itemsize - last_offset + if padding > 0: + _fields_.append(('', ctypes.c_char * padding)) + + # we inserted manual padding, so always `_pack_` + return type('struct', (ctypes.Structure,), dict( + _fields_=_fields_, + _pack_=1, + __module__=None, + )) + + + def _ctype_from_dtype(dtype): + if dtype.fields is not None: + return _ctype_from_dtype_structured(dtype) + elif dtype.subdtype is not None: + return _ctype_from_dtype_subarray(dtype) + else: + return _ctype_from_dtype_scalar(dtype) + + + def as_ctypes_type(dtype): + r""" + Convert a dtype into a ctypes type. + + Parameters + ---------- + dtype : dtype + The dtype to convert + + Returns + ------- + ctype + A ctype scalar, union, array, or struct + + Raises + ------ + NotImplementedError + If the conversion is not possible + + Notes + ----- + This function does not losslessly round-trip in either direction. + + ``np.dtype(as_ctypes_type(dt))`` will: + + - insert padding fields + - reorder fields to be sorted by offset + - discard field titles + + ``as_ctypes_type(np.dtype(ctype))`` will: + + - discard the class names of `ctypes.Structure`\ s and + `ctypes.Union`\ s + - convert single-element `ctypes.Union`\ s into single-element + `ctypes.Structure`\ s + - insert padding fields + + """ + return _ctype_from_dtype(_dtype(dtype)) + + + def as_array(obj, shape=None): + """ + Create a numpy array from a ctypes array or POINTER. + + The numpy array shares the memory with the ctypes object. + + The shape parameter must be given if converting from a ctypes POINTER. + The shape parameter is ignored if converting from a ctypes array + """ + if isinstance(obj, ctypes._Pointer): + # convert pointers to an array of the desired shape + if shape is None: + raise TypeError( + 'as_array() requires a shape argument when called on a ' + 'pointer') + p_arr_type = ctypes.POINTER(_ctype_ndarray(obj._type_, shape)) + obj = ctypes.cast(obj, p_arr_type).contents + + return asarray(obj) + + + def as_ctypes(obj): + """Create and return a ctypes object from a numpy array. Actually + anything that exposes the __array_interface__ is accepted.""" + ai = obj.__array_interface__ + if ai["strides"]: + raise TypeError("strided arrays not supported") + if ai["version"] != 3: + raise TypeError("only __array_interface__ version 3 supported") + addr, readonly = ai["data"] + if readonly: + raise TypeError("readonly arrays unsupported") + + # can't use `_dtype((ai["typestr"], ai["shape"]))` here, as it overflows + # dtype.itemsize (gh-14214) + ctype_scalar = as_ctypes_type(ai["typestr"]) + result_type = _ctype_ndarray(ctype_scalar, ai["shape"]) + result = result_type.from_address(addr) + result.__keep = obj + return result diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/exceptions.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..2f843810141a7c2d78de9ff75f5a0db9e592c981 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/exceptions.py @@ -0,0 +1,231 @@ +""" +Exceptions and Warnings (:mod:`numpy.exceptions`) +================================================= + +General exceptions used by NumPy. Note that some exceptions may be module +specific, such as linear algebra errors. + +.. versionadded:: NumPy 1.25 + + The exceptions module is new in NumPy 1.25. Older exceptions remain + available through the main NumPy namespace for compatibility. + +.. currentmodule:: numpy.exceptions + +Warnings +-------- +.. autosummary:: + :toctree: generated/ + + ComplexWarning Given when converting complex to real. + VisibleDeprecationWarning Same as a DeprecationWarning, but more visible. + +Exceptions +---------- +.. autosummary:: + :toctree: generated/ + + AxisError Given when an axis was invalid. + DTypePromotionError Given when no common dtype could be found. + TooHardError Error specific to `numpy.shares_memory`. + +""" + + +__all__ = [ + "ComplexWarning", "VisibleDeprecationWarning", "ModuleDeprecationWarning", + "TooHardError", "AxisError", "DTypePromotionError"] + + +# Disallow reloading this module so as to preserve the identities of the +# classes defined here. +if '_is_loaded' in globals(): + raise RuntimeError('Reloading numpy._globals is not allowed') +_is_loaded = True + + +class ComplexWarning(RuntimeWarning): + """ + The warning raised when casting a complex dtype to a real dtype. + + As implemented, casting a complex number to a real discards its imaginary + part, but this behavior may not be what the user actually wants. + + """ + pass + + +class ModuleDeprecationWarning(DeprecationWarning): + """Module deprecation warning. + + .. warning:: + + This warning should not be used, since nose testing is not relevant + anymore. + + The nose tester turns ordinary Deprecation warnings into test failures. + That makes it hard to deprecate whole modules, because they get + imported by default. So this is a special Deprecation warning that the + nose tester will let pass without making tests fail. + + """ + + +class VisibleDeprecationWarning(UserWarning): + """Visible deprecation warning. + + By default, python will not show deprecation warnings, so this class + can be used when a very visible warning is helpful, for example because + the usage is most likely a user bug. + + """ + + +# Exception used in shares_memory() +class TooHardError(RuntimeError): + """max_work was exceeded. + + This is raised whenever the maximum number of candidate solutions + to consider specified by the ``max_work`` parameter is exceeded. + Assigning a finite number to max_work may have caused the operation + to fail. + + """ + + pass + + +class AxisError(ValueError, IndexError): + """Axis supplied was invalid. + + This is raised whenever an ``axis`` parameter is specified that is larger + than the number of array dimensions. + For compatibility with code written against older numpy versions, which + raised a mixture of `ValueError` and `IndexError` for this situation, this + exception subclasses both to ensure that ``except ValueError`` and + ``except IndexError`` statements continue to catch `AxisError`. + + .. versionadded:: 1.13 + + Parameters + ---------- + axis : int or str + The out of bounds axis or a custom exception message. + If an axis is provided, then `ndim` should be specified as well. + ndim : int, optional + The number of array dimensions. + msg_prefix : str, optional + A prefix for the exception message. + + Attributes + ---------- + axis : int, optional + The out of bounds axis or ``None`` if a custom exception + message was provided. This should be the axis as passed by + the user, before any normalization to resolve negative indices. + + .. versionadded:: 1.22 + ndim : int, optional + The number of array dimensions or ``None`` if a custom exception + message was provided. + + .. versionadded:: 1.22 + + + Examples + -------- + >>> array_1d = np.arange(10) + >>> np.cumsum(array_1d, axis=1) + Traceback (most recent call last): + ... + numpy.exceptions.AxisError: axis 1 is out of bounds for array of dimension 1 + + Negative axes are preserved: + + >>> np.cumsum(array_1d, axis=-2) + Traceback (most recent call last): + ... + numpy.exceptions.AxisError: axis -2 is out of bounds for array of dimension 1 + + The class constructor generally takes the axis and arrays' + dimensionality as arguments: + + >>> print(np.AxisError(2, 1, msg_prefix='error')) + error: axis 2 is out of bounds for array of dimension 1 + + Alternatively, a custom exception message can be passed: + + >>> print(np.AxisError('Custom error message')) + Custom error message + + """ + + __slots__ = ("axis", "ndim", "_msg") + + def __init__(self, axis, ndim=None, msg_prefix=None): + if ndim is msg_prefix is None: + # single-argument form: directly set the error message + self._msg = axis + self.axis = None + self.ndim = None + else: + self._msg = msg_prefix + self.axis = axis + self.ndim = ndim + + def __str__(self): + axis = self.axis + ndim = self.ndim + + if axis is ndim is None: + return self._msg + else: + msg = f"axis {axis} is out of bounds for array of dimension {ndim}" + if self._msg is not None: + msg = f"{self._msg}: {msg}" + return msg + + +class DTypePromotionError(TypeError): + """Multiple DTypes could not be converted to a common one. + + This exception derives from ``TypeError`` and is raised whenever dtypes + cannot be converted to a single common one. This can be because they + are of a different category/class or incompatible instances of the same + one (see Examples). + + Notes + ----- + Many functions will use promotion to find the correct result and + implementation. For these functions the error will typically be chained + with a more specific error indicating that no implementation was found + for the input dtypes. + + Typically promotion should be considered "invalid" between the dtypes of + two arrays when `arr1 == arr2` can safely return all ``False`` because the + dtypes are fundamentally different. + + Examples + -------- + Datetimes and complex numbers are incompatible classes and cannot be + promoted: + + >>> np.result_type(np.dtype("M8[s]"), np.complex128) + DTypePromotionError: The DType could not + be promoted by . This means that no common + DType exists for the given inputs. For example they cannot be stored in a + single array unless the dtype is `object`. The full list of DTypes is: + (, ) + + For example for structured dtypes, the structure can mismatch and the + same ``DTypePromotionError`` is given when two structured dtypes with + a mismatch in their number of fields is given: + + >>> dtype1 = np.dtype([("field1", np.float64), ("field2", np.int64)]) + >>> dtype2 = np.dtype([("field1", np.float64)]) + >>> np.promote_types(dtype1, dtype2) + DTypePromotionError: field names `('field1', 'field2')` and `('field1',)` + mismatch. + + """ + pass diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/matlib.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/matlib.py new file mode 100644 index 0000000000000000000000000000000000000000..e929fd9b1885f208afb6301f19cc21511adc098b --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/numpy/matlib.py @@ -0,0 +1,378 @@ +import warnings + +# 2018-05-29, PendingDeprecationWarning added to matrix.__new__ +# 2020-01-23, numpy 1.19.0 PendingDeprecatonWarning +warnings.warn("Importing from numpy.matlib is deprecated since 1.19.0. " + "The matrix subclass is not the recommended way to represent " + "matrices or deal with linear algebra (see " + "https://docs.scipy.org/doc/numpy/user/numpy-for-matlab-users.html). " + "Please adjust your code to use regular ndarray. ", + PendingDeprecationWarning, stacklevel=2) + +import numpy as np +from numpy.matrixlib.defmatrix import matrix, asmatrix +# Matlib.py contains all functions in the numpy namespace with a few +# replacements. See doc/source/reference/routines.matlib.rst for details. +# Need * as we're copying the numpy namespace. +from numpy import * # noqa: F403 + +__version__ = np.__version__ + +__all__ = np.__all__[:] # copy numpy namespace +__all__ += ['rand', 'randn', 'repmat'] + +def empty(shape, dtype=None, order='C'): + """Return a new matrix of given shape and type, without initializing entries. + + Parameters + ---------- + shape : int or tuple of int + Shape of the empty matrix. + dtype : data-type, optional + Desired output data-type. + order : {'C', 'F'}, optional + Whether to store multi-dimensional data in row-major + (C-style) or column-major (Fortran-style) order in + memory. + + See Also + -------- + empty_like, zeros + + Notes + ----- + `empty`, unlike `zeros`, does not set the matrix values to zero, + and may therefore be marginally faster. On the other hand, it requires + the user to manually set all the values in the array, and should be + used with caution. + + Examples + -------- + >>> import numpy.matlib + >>> np.matlib.empty((2, 2)) # filled with random data + matrix([[ 6.76425276e-320, 9.79033856e-307], # random + [ 7.39337286e-309, 3.22135945e-309]]) + >>> np.matlib.empty((2, 2), dtype=int) + matrix([[ 6600475, 0], # random + [ 6586976, 22740995]]) + + """ + return ndarray.__new__(matrix, shape, dtype, order=order) + +def ones(shape, dtype=None, order='C'): + """ + Matrix of ones. + + Return a matrix of given shape and type, filled with ones. + + Parameters + ---------- + shape : {sequence of ints, int} + Shape of the matrix + dtype : data-type, optional + The desired data-type for the matrix, default is np.float64. + order : {'C', 'F'}, optional + Whether to store matrix in C- or Fortran-contiguous order, + default is 'C'. + + Returns + ------- + out : matrix + Matrix of ones of given shape, dtype, and order. + + See Also + -------- + ones : Array of ones. + matlib.zeros : Zero matrix. + + Notes + ----- + If `shape` has length one i.e. ``(N,)``, or is a scalar ``N``, + `out` becomes a single row matrix of shape ``(1,N)``. + + Examples + -------- + >>> np.matlib.ones((2,3)) + matrix([[1., 1., 1.], + [1., 1., 1.]]) + + >>> np.matlib.ones(2) + matrix([[1., 1.]]) + + """ + a = ndarray.__new__(matrix, shape, dtype, order=order) + a.fill(1) + return a + +def zeros(shape, dtype=None, order='C'): + """ + Return a matrix of given shape and type, filled with zeros. + + Parameters + ---------- + shape : int or sequence of ints + Shape of the matrix + dtype : data-type, optional + The desired data-type for the matrix, default is float. + order : {'C', 'F'}, optional + Whether to store the result in C- or Fortran-contiguous order, + default is 'C'. + + Returns + ------- + out : matrix + Zero matrix of given shape, dtype, and order. + + See Also + -------- + numpy.zeros : Equivalent array function. + matlib.ones : Return a matrix of ones. + + Notes + ----- + If `shape` has length one i.e. ``(N,)``, or is a scalar ``N``, + `out` becomes a single row matrix of shape ``(1,N)``. + + Examples + -------- + >>> import numpy.matlib + >>> np.matlib.zeros((2, 3)) + matrix([[0., 0., 0.], + [0., 0., 0.]]) + + >>> np.matlib.zeros(2) + matrix([[0., 0.]]) + + """ + a = ndarray.__new__(matrix, shape, dtype, order=order) + a.fill(0) + return a + +def identity(n,dtype=None): + """ + Returns the square identity matrix of given size. + + Parameters + ---------- + n : int + Size of the returned identity matrix. + dtype : data-type, optional + Data-type of the output. Defaults to ``float``. + + Returns + ------- + out : matrix + `n` x `n` matrix with its main diagonal set to one, + and all other elements zero. + + See Also + -------- + numpy.identity : Equivalent array function. + matlib.eye : More general matrix identity function. + + Examples + -------- + >>> import numpy.matlib + >>> np.matlib.identity(3, dtype=int) + matrix([[1, 0, 0], + [0, 1, 0], + [0, 0, 1]]) + + """ + a = array([1]+n*[0], dtype=dtype) + b = empty((n, n), dtype=dtype) + b.flat = a + return b + +def eye(n,M=None, k=0, dtype=float, order='C'): + """ + Return a matrix with ones on the diagonal and zeros elsewhere. + + Parameters + ---------- + n : int + Number of rows in the output. + M : int, optional + Number of columns in the output, defaults to `n`. + k : int, optional + Index of the diagonal: 0 refers to the main diagonal, + a positive value refers to an upper diagonal, + and a negative value to a lower diagonal. + dtype : dtype, optional + Data-type of the returned matrix. + order : {'C', 'F'}, optional + Whether the output should be stored in row-major (C-style) or + column-major (Fortran-style) order in memory. + + .. versionadded:: 1.14.0 + + Returns + ------- + I : matrix + A `n` x `M` matrix where all elements are equal to zero, + except for the `k`-th diagonal, whose values are equal to one. + + See Also + -------- + numpy.eye : Equivalent array function. + identity : Square identity matrix. + + Examples + -------- + >>> import numpy.matlib + >>> np.matlib.eye(3, k=1, dtype=float) + matrix([[0., 1., 0.], + [0., 0., 1.], + [0., 0., 0.]]) + + """ + return asmatrix(np.eye(n, M=M, k=k, dtype=dtype, order=order)) + +def rand(*args): + """ + Return a matrix of random values with given shape. + + Create a matrix of the given shape and propagate it with + random samples from a uniform distribution over ``[0, 1)``. + + Parameters + ---------- + \\*args : Arguments + Shape of the output. + If given as N integers, each integer specifies the size of one + dimension. + If given as a tuple, this tuple gives the complete shape. + + Returns + ------- + out : ndarray + The matrix of random values with shape given by `\\*args`. + + See Also + -------- + randn, numpy.random.RandomState.rand + + Examples + -------- + >>> np.random.seed(123) + >>> import numpy.matlib + >>> np.matlib.rand(2, 3) + matrix([[0.69646919, 0.28613933, 0.22685145], + [0.55131477, 0.71946897, 0.42310646]]) + >>> np.matlib.rand((2, 3)) + matrix([[0.9807642 , 0.68482974, 0.4809319 ], + [0.39211752, 0.34317802, 0.72904971]]) + + If the first argument is a tuple, other arguments are ignored: + + >>> np.matlib.rand((2, 3), 4) + matrix([[0.43857224, 0.0596779 , 0.39804426], + [0.73799541, 0.18249173, 0.17545176]]) + + """ + if isinstance(args[0], tuple): + args = args[0] + return asmatrix(np.random.rand(*args)) + +def randn(*args): + """ + Return a random matrix with data from the "standard normal" distribution. + + `randn` generates a matrix filled with random floats sampled from a + univariate "normal" (Gaussian) distribution of mean 0 and variance 1. + + Parameters + ---------- + \\*args : Arguments + Shape of the output. + If given as N integers, each integer specifies the size of one + dimension. If given as a tuple, this tuple gives the complete shape. + + Returns + ------- + Z : matrix of floats + A matrix of floating-point samples drawn from the standard normal + distribution. + + See Also + -------- + rand, numpy.random.RandomState.randn + + Notes + ----- + For random samples from the normal distribution with mean ``mu`` and + standard deviation ``sigma``, use:: + + sigma * np.matlib.randn(...) + mu + + Examples + -------- + >>> np.random.seed(123) + >>> import numpy.matlib + >>> np.matlib.randn(1) + matrix([[-1.0856306]]) + >>> np.matlib.randn(1, 2, 3) + matrix([[ 0.99734545, 0.2829785 , -1.50629471], + [-0.57860025, 1.65143654, -2.42667924]]) + + Two-by-four matrix of samples from the normal distribution with + mean 3 and standard deviation 2.5: + + >>> 2.5 * np.matlib.randn((2, 4)) + 3 + matrix([[1.92771843, 6.16484065, 0.83314899, 1.30278462], + [2.76322758, 6.72847407, 1.40274501, 1.8900451 ]]) + + """ + if isinstance(args[0], tuple): + args = args[0] + return asmatrix(np.random.randn(*args)) + +def repmat(a, m, n): + """ + Repeat a 0-D to 2-D array or matrix MxN times. + + Parameters + ---------- + a : array_like + The array or matrix to be repeated. + m, n : int + The number of times `a` is repeated along the first and second axes. + + Returns + ------- + out : ndarray + The result of repeating `a`. + + Examples + -------- + >>> import numpy.matlib + >>> a0 = np.array(1) + >>> np.matlib.repmat(a0, 2, 3) + array([[1, 1, 1], + [1, 1, 1]]) + + >>> a1 = np.arange(4) + >>> np.matlib.repmat(a1, 2, 2) + array([[0, 1, 2, 3, 0, 1, 2, 3], + [0, 1, 2, 3, 0, 1, 2, 3]]) + + >>> a2 = np.asmatrix(np.arange(6).reshape(2, 3)) + >>> np.matlib.repmat(a2, 2, 3) + matrix([[0, 1, 2, 0, 1, 2, 0, 1, 2], + [3, 4, 5, 3, 4, 5, 3, 4, 5], + [0, 1, 2, 0, 1, 2, 0, 1, 2], + [3, 4, 5, 3, 4, 5, 3, 4, 5]]) + + """ + a = asanyarray(a) + ndim = a.ndim + if ndim == 0: + origrows, origcols = (1, 1) + elif ndim == 1: + origrows, origcols = (1, a.shape[0]) + else: + origrows, origcols = a.shape + rows = origrows * m + cols = origcols * n + c = a.reshape(1, a.size).repeat(m, 0).reshape(rows, origcols).repeat(n, 0) + return c.reshape(rows, cols) diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..336db3773f762d840c32afe417e0f26e817f315e --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/__init__.py @@ -0,0 +1,324 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_torch_greater_or_equal + + +_import_structure = { + "aqlm": ["replace_with_aqlm_linear"], + "awq": [ + "post_init_awq_exllama_modules", + "replace_quantization_scales", + "replace_with_awq_linear", + ], + "bitnet": [ + "BitLinear", + "pack_weights", + "replace_with_bitnet_linear", + "unpack_weights", + ], + "bitsandbytes": [ + "Bnb4bitQuantize", + "dequantize_and_replace", + "replace_with_bnb_linear", + "validate_bnb_backend_availability", + ], + "deepspeed": [ + "HfDeepSpeedConfig", + "HfTrainerDeepSpeedConfig", + "deepspeed_config", + "deepspeed_init", + "deepspeed_load_checkpoint", + "deepspeed_optim_sched", + "is_deepspeed_available", + "is_deepspeed_zero3_enabled", + "set_hf_deepspeed_config", + "unset_hf_deepspeed_config", + ], + "eetq": ["replace_with_eetq_linear"], + "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"], + "finegrained_fp8": ["FP8Linear", "replace_with_fp8_linear"], + "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], + "ggml": [ + "GGUF_CONFIG_DEFAULTS_MAPPING", + "GGUF_CONFIG_MAPPING", + "GGUF_TOKENIZER_MAPPING", + "_gguf_parse_value", + "load_dequant_gguf_tensor", + "load_gguf", + ], + "higgs": [ + "HiggsLinear", + "dequantize_higgs", + "quantize_with_higgs", + "replace_with_higgs_linear", + ], + "hqq": ["prepare_for_hqq_linear"], + "hub_kernels": [ + "LayerRepository", + "lazy_load_kernel", + "register_kernel_mapping", + "replace_kernel_forward_from_hub", + "use_kernel_forward_from_hub", + "use_kernel_func_from_hub", + "use_kernelized_func", + ], + "integration_utils": [ + "INTEGRATION_TO_CALLBACK", + "AzureMLCallback", + "ClearMLCallback", + "CodeCarbonCallback", + "CometCallback", + "DagsHubCallback", + "DVCLiveCallback", + "FlyteCallback", + "KubeflowCallback", + "MLflowCallback", + "NeptuneCallback", + "NeptuneMissingConfiguration", + "SwanLabCallback", + "TensorBoardCallback", + "TrackioCallback", + "WandbCallback", + "get_available_reporting_integrations", + "get_reporting_integration_callbacks", + "hp_params", + "is_azureml_available", + "is_clearml_available", + "is_codecarbon_available", + "is_comet_available", + "is_dagshub_available", + "is_dvclive_available", + "is_flyte_deck_standard_available", + "is_flytekit_available", + "is_kubeflow_available", + "is_mlflow_available", + "is_neptune_available", + "is_optuna_available", + "is_ray_available", + "is_ray_tune_available", + "is_swanlab_available", + "is_tensorboard_available", + "is_trackio_available", + "is_wandb_available", + "rewrite_logs", + "run_hp_search_optuna", + "run_hp_search_ray", + "run_hp_search_wandb", + ], + "liger": ["apply_liger_kernel"], + "metal_quantization": [ + "MetalLinear", + "replace_with_metal_linear", + ], + "moe": [ + "batched_mm_experts_forward", + "grouped_mm_experts_forward", + "use_experts_implementation", + ], + "mxfp4": [ + "Mxfp4GptOssExperts", + "convert_moe_packed_tensors", + "dequantize", + "load_and_swizzle_mxfp4", + "quantize_to_mxfp4", + "replace_with_mxfp4_linear", + "swizzle_mxfp4", + ], + "neftune": [ + "activate_neftune", + "deactivate_neftune", + "neftune_post_forward_hook", + ], + "peft": ["PeftAdapterMixin"], + "quanto": ["replace_with_quanto_layers"], + "sinq": ["SinqDeserialize", "SinqQuantize"], + "spqr": ["replace_with_spqr_linear"], + "vptq": ["replace_with_vptq_linear"], +} + +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["executorch"] = [ + "TorchExportableModuleWithStaticCache", + "convert_and_export_with_cache", + ] + +_import_structure["tensor_parallel"] = [ + "shard_and_distribute_module", + "ALL_PARALLEL_STYLES", + "translate_to_torch_parallel_style", +] +try: + if not is_torch_greater_or_equal("2.5"): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["flex_attention"] = [ + "make_flex_block_causal_mask", + ] + +if TYPE_CHECKING: + from .aqlm import replace_with_aqlm_linear + from .awq import ( + post_init_awq_exllama_modules, + replace_quantization_scales, + replace_with_awq_linear, + ) + from .bitnet import ( + BitLinear, + pack_weights, + replace_with_bitnet_linear, + unpack_weights, + ) + from .bitsandbytes import ( + Bnb4bitQuantize, + dequantize_and_replace, + replace_with_bnb_linear, + validate_bnb_backend_availability, + ) + from .deepspeed import ( + HfDeepSpeedConfig, + HfTrainerDeepSpeedConfig, + deepspeed_config, + deepspeed_init, + deepspeed_load_checkpoint, + deepspeed_optim_sched, + is_deepspeed_available, + is_deepspeed_zero3_enabled, + set_hf_deepspeed_config, + unset_hf_deepspeed_config, + ) + from .eetq import replace_with_eetq_linear + from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear + from .finegrained_fp8 import FP8Linear, replace_with_fp8_linear + from .fsdp import is_fsdp_enabled, is_fsdp_managed_module + from .ggml import ( + GGUF_CONFIG_DEFAULTS_MAPPING, + GGUF_CONFIG_MAPPING, + GGUF_TOKENIZER_MAPPING, + _gguf_parse_value, + load_dequant_gguf_tensor, + load_gguf, + ) + from .higgs import HiggsLinear, dequantize_higgs, quantize_with_higgs, replace_with_higgs_linear + from .hqq import prepare_for_hqq_linear + from .hub_kernels import ( + LayerRepository, + lazy_load_kernel, + register_kernel_mapping, + replace_kernel_forward_from_hub, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, + ) + from .integration_utils import ( + INTEGRATION_TO_CALLBACK, + AzureMLCallback, + ClearMLCallback, + CodeCarbonCallback, + CometCallback, + DagsHubCallback, + DVCLiveCallback, + FlyteCallback, + KubeflowCallback, + MLflowCallback, + NeptuneCallback, + NeptuneMissingConfiguration, + SwanLabCallback, + TensorBoardCallback, + TrackioCallback, + WandbCallback, + get_available_reporting_integrations, + get_reporting_integration_callbacks, + hp_params, + is_azureml_available, + is_clearml_available, + is_codecarbon_available, + is_comet_available, + is_dagshub_available, + is_dvclive_available, + is_flyte_deck_standard_available, + is_flytekit_available, + is_kubeflow_available, + is_mlflow_available, + is_neptune_available, + is_optuna_available, + is_ray_available, + is_ray_tune_available, + is_swanlab_available, + is_tensorboard_available, + is_trackio_available, + is_wandb_available, + rewrite_logs, + run_hp_search_optuna, + run_hp_search_ray, + run_hp_search_wandb, + ) + from .liger import apply_liger_kernel + from .metal_quantization import ( + MetalLinear, + replace_with_metal_linear, + ) + from .moe import ( + batched_mm_experts_forward, + grouped_mm_experts_forward, + use_experts_implementation, + ) + from .mxfp4 import ( + Mxfp4GptOssExperts, + dequantize, + load_and_swizzle_mxfp4, + quantize_to_mxfp4, + replace_with_mxfp4_linear, + swizzle_mxfp4, + ) + from .neftune import activate_neftune, deactivate_neftune, neftune_post_forward_hook + from .peft import PeftAdapterMixin + from .quanto import replace_with_quanto_layers + from .sinq import SinqDeserialize, SinqQuantize + from .spqr import replace_with_spqr_linear + from .vptq import replace_with_vptq_linear + + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .executorch import TorchExportableModuleWithStaticCache, convert_and_export_with_cache + + from .tensor_parallel import ( + ALL_PARALLEL_STYLES, + shard_and_distribute_module, + translate_to_torch_parallel_style, + ) + + try: + if not is_torch_greater_or_equal("2.5"): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .flex_attention import make_flex_block_causal_mask +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/awq.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/awq.py new file mode 100644 index 0000000000000000000000000000000000000000..059343ea5c78dc43cfd03924dc27d9a497a464e9 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/awq.py @@ -0,0 +1,119 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"AWQ (Activation aware Weight Quantization) integration file" + +from ..quantizers.quantizers_utils import should_convert_module +from ..utils import is_torch_available, logging + + +if is_torch_available(): + import torch + import torch.nn as nn + +logger = logging.get_logger(__name__) + + +AWQ_SCALES_MAPPINGS = { + "starcoder2": {"act": "act", "layer_before_act": "c_fc"}, + "RefinedWebModel": {"act": "act", "layer_before_act": "dense_h_to_4h"}, + "falcon": {"act": "act", "layer_before_act": "dense_h_to_4h"}, + "mpt": {"act": "act", "layer_before_act": "up_proj"}, + "gptj": {"act": "act", "layer_before_act": "fc_in"}, + "gpt_neox": {"act": "act", "layer_before_act": "dense_h_to_4h"}, + "gpt_bigcode": {"act": "act", "layer_before_act": "c_fc"}, + "bloom": {"act": "gelu_impl", "layer_before_act": "dense_h_to_4h"}, +} + + +def replace_quantization_scales(model, model_type): + from gptqmodel.quantization.awq.modules.act import ScaledActivation + + if model_type not in AWQ_SCALES_MAPPINGS: + return model + for name, module in model.named_children(): + act_name = AWQ_SCALES_MAPPINGS[model_type]["act"] + layer_before_act_name = AWQ_SCALES_MAPPINGS[model_type]["layer_before_act"] + if name == act_name and hasattr(model, layer_before_act_name): + layer_before_act = getattr(model, AWQ_SCALES_MAPPINGS[model_type]["layer_before_act"]) + size = layer_before_act.out_features + scale_like = torch.ones(size) + model._modules[name] = ScaledActivation(module, scale_like) + _ = replace_quantization_scales(module, model_type) + return model + + +def replace_with_awq_linear( + model, + modules_to_not_convert=None, + quantization_config=None, + device_map: str | dict | None = None, +) -> bool: + """ + Public method that replaces the linear layers of the given model with awq quantized layers. + + Args: + model (`torch.nn.Module`): + The model to convert, can be any `torch.nn.Module` instance. + quantization_config (`AwqConfig`): + The quantization config object that contains the quantization parameters. + modules_to_not_convert (`list[str]`, *optional*, defaults to `None`): + A list of nn.Linear weights to not convert. If a parameter path is in the list (e.g. `lm_head.weight`), the corresponding module will not be + converted. + device_map (`Union[str, dict]`, *optional*, defaults to `None`): + The device map that maps the parameters to the device + """ + from gptqmodel.quantization import METHOD + from gptqmodel.utils.importer import hf_select_quant_linear_v2 + + target_cls = hf_select_quant_linear_v2( + bits=quantization_config.bits, + group_size=quantization_config.group_size, + desc_act=False, + sym=False, + format=quantization_config.format, + backend=quantization_config.backend, + device_map=device_map, + quant_method=METHOD.AWQ, + zero_point=quantization_config.zero_point, + pack=False, + ) + + for module_name, module in model.named_modules(): + if not should_convert_module(module_name, modules_to_not_convert): + continue + with torch.device("meta"): + if isinstance(module, nn.Linear): + new_module = target_cls( + bits=quantization_config.bits, + sym=quantization_config.sym, + desc_act=quantization_config.desc_act, + group_size=quantization_config.group_size, + in_features=module.in_features, + out_features=module.out_features, + bias=module.bias is not None, + dev=module.weight.device, + register_buffers=True, + ) + new_module.requires_grad_(False) + model.set_submodule(module_name, new_module) + has_been_replaced = True + + if not has_been_replaced: + logger.warning( + "You are loading your model using eetq but no linear modules were found in your model." + " Please double check your model architecture, or submit an issue on github if you think this is" + " a bug." + ) + + return model diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/executorch.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/executorch.py new file mode 100644 index 0000000000000000000000000000000000000000..675a0ea5783a71fac1c03495825240aa5993c1b1 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/executorch.py @@ -0,0 +1,1137 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on +# an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the +# specific language governing permissions and limitations under the License. + +import logging + +import torch + +from ..cache_utils import ( + DynamicCache, + DynamicLayer, + DynamicSlidingWindowLayer, + EncoderDecoderCache, + StaticCache, + StaticLayer, + StaticSlidingWindowLayer, +) +from ..generation.configuration_utils import GenerationConfig +from ..modeling_utils import PreTrainedModel +from ..pytorch_utils import ( + is_torch_greater_or_equal, + is_torch_greater_or_equal_than_2_6, +) + + +class TorchExportableModuleForVLM: + """ + A wrapper class for exporting Vision-Language Models (VLMs) like SmolVLM2 for ExecuTorch. + + This class handles the export of three main components: + 1. Vision encoder (processes images to visual features) + 2. Connector/projector (maps visual features to text embedding space) + 3. Text decoder (generates text from combined visual and text tokens) + """ + + def __init__(self, model, max_batch_size: int = 1, max_cache_len: int = 1024): + """ + Initialize the exportable VLM module. + + Args: + model: The VLM (e.g. SmolVLM) model instance + max_batch_size: Maximum batch size. Always 1 for ExecuTorch + max_cache_len: Maximum cache length for text generation + """ + self.model = model + self.max_batch_size = max_batch_size + self.max_cache_len = max_cache_len + self.config = model.config + + # Extract individual components + self.vision_encoder = model.model.vision_model + self.connector = model.model.connector + self.text_decoder = model.model.text_model + + # Store exported programs + self.exported_vision_encoder = None + self.exported_connector = None + self.exported_text_decoder = None + + def export_vision_encoder(self): + """Export the vision encoder component.""" + self.vision_encoder.eval() + + # Create example input + pixel_values = torch.randn(1, 3, 384, 384, dtype=torch.float32) + + # Define dynamic shapes + dynamic_shapes = { + "pixel_values": { + 2: torch.export.Dim.AUTO, + 3: torch.export.Dim.AUTO, + } + } + + self.exported_vision_encoder = torch.export.export( + self.vision_encoder, + args=(pixel_values,), + dynamic_shapes=dynamic_shapes, + strict=False, + ) + + return self.exported_vision_encoder + + def export_connector(self): + """Export the connector component.""" + self.connector.eval() + + # Vision encoder output shape: [batch_size, num_patches, vision_hidden_size] + vision_hidden_size = self.config.vision_config.hidden_size + image_size = self.config.vision_config.image_size + patch_size = self.config.vision_config.patch_size + patches_per_dim = image_size // patch_size + num_patches = patches_per_dim * patches_per_dim + image_hidden_states = torch.randn(1, num_patches, vision_hidden_size, dtype=torch.float32) + + # Define dynamic shapes - static batch_size=1, dynamic num_patches + dynamic_shapes = {"image_hidden_states": {1: torch.export.Dim.AUTO}} + + # Export the connector using torch.export + self.exported_connector = torch.export.export( + self.connector, + args=(image_hidden_states,), + dynamic_shapes=dynamic_shapes, + strict=False, + ) + + return self.exported_connector + + def export_text_decoder(self): + """Export the text decoder component.""" + + # Create text decoder exportable wrapper + self.exportable_text_decoder = TorchExportableModuleForDecoderOnlyLM(model=self.text_decoder) + + # Use the existing text decoder exportable wrapper + seq_length = 3 + input_ids = torch.zeros((1, seq_length), dtype=torch.long) + cache_position = torch.arange(seq_length, dtype=torch.long) + max_seq_length = min(self.max_cache_len, self.config.text_config.max_position_embeddings) + seq_len_dim = torch.export.Dim("seq_length_dim", max=max_seq_length - 1) + + dynamic_shapes = { + "input_ids": {1: seq_len_dim}, + "cache_position": {0: seq_len_dim}, + } + + self.exported_text_decoder = self.exportable_text_decoder.export( + input_ids=input_ids, + cache_position=cache_position, + dynamic_shapes=dynamic_shapes, + strict=False, + ) + + return self.exported_text_decoder + + def export(self, **kwargs): + """Export all components of the VLM model.""" + self.export_vision_encoder(**kwargs) + self.export_connector(**kwargs) + self.export_text_decoder(**kwargs) + return { + "vision_encoder": self.exported_vision_encoder, + "connector": self.exported_connector, + "text_decoder": self.exported_text_decoder, + } + + def forward(self, pixel_values, input_ids, cache_position): + """ + Simplified forward pass for inference with guaranteed non-null input_ids and cache_position. + + Args: + pixel_values: Input images [1, channels, height, width] (optional) + input_ids: Text token IDs [1, seq_len] (required - won't be None) + cache_position: Cache positions [seq_len] (required - won't be None) + + Returns: + Output with logits for text generation + """ + + def generate( + self, pixel_values=None, input_ids=None, max_new_tokens=50, do_sample=False, temperature=1.0, **kwargs + ): + """ + Simplified generate method with guaranteed non-null input_ids. + + Args: + pixel_values: Input images [1, channels, height, width] (optional) + input_ids: Initial text tokens [1, seq_len] (required - won't be None) + max_new_tokens: Maximum number of tokens to generate + do_sample: Whether to use sampling or greedy decoding + temperature: Temperature for sampling + + Returns: + Generated sequences + """ + + +class TorchExportableModuleForDecoderOnlyLM(torch.nn.Module): + """ + A recipe module designed to make a `PreTrainedModel` exportable with `torch.export`, + specifically for decoder-only LM with cache. This module ensures that the + exported model is compatible with further lowering and execution in `ExecuTorch`. + """ + + def __init__( + self, + model: PreTrainedModel, + batch_size: int | None = None, + max_cache_len: int | None = None, + device: torch.device | None = None, + ) -> None: + """ + Initializes the exportable module. + + Args: + model (`PreTrainedModel`): The pretrained model to wrap. + + Raises: + ValueError: If the model is configured with a unsupported cache implementation. + """ + super().__init__() + + config = model.config.get_text_config() + + if not hasattr(config, "use_cache") or config.use_cache is False: + raise ValueError("The model must have caching enabled to be performant.") + + if hasattr(config, "layer_types") and getattr(config, "sliding_window", None) is not None: + self.model = TorchExportableModuleWithHybridCache(model, batch_size, max_cache_len, device) + else: + # If `layer_types` is not specified explicitly in the config or `sliding_window` is null, + # there is only 1 type of layers, so export will use `StaticCache` by default. + logging.info( + "Using `StaticCache` for export as `layer_types` is not specified or `sliding_window` is `null` in the config." + ) + self.model = TorchExportableModuleWithStaticCache(model, batch_size, max_cache_len, device) + + def forward( + self, + input_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + cache_position: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + Forward pass of the module, which is compatible with the ExecuTorch llm runner. + + Args: + input_ids (`torch.Tensor`): Tensor representing current input token id to the module. + inputs_embeds (`torch.Tensor`): Tensor representing current input embeddings to the module. + cache_position (`torch.Tensor`): Tensor representing current input position in the cache. + + Returns: + torch.Tensor: Logits output from the model. + """ + return self.model.forward(input_ids=input_ids, inputs_embeds=inputs_embeds) + + def export( + self, + input_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + cache_position: torch.Tensor | None = None, + dynamic_shapes: dict | None = None, + strict: bool | None = None, + ) -> torch.export.ExportedProgram: + """ + Export the wrapped module using `torch.export`. + + Args: + input_ids (`Optional[torch.Tensor]`): + Tensor representing current input token id to the module. Must specify either this or inputs_embeds. + inputs_embeds (`Optional[torch.Tensor]`): + Tensor representing current input embeddings to the module. Must specify either this or input_ids. + cache_position (`Optional[torch.Tensor]`): + Tensor representing current input position in the cache. If not provided, a default tensor will be used. + dynamic_shapes (`Optional[dict]`): + Dynamic shapes to use for export if specified. + strict(`Optional[bool]`): + Flag to instruct `torch.export` to use `torchdynamo`. + + Returns: + torch.export.ExportedProgram: The exported program that can be used for inference. + + Examples: + Export with input_ids: + ```python + # Prepare inputs + input_ids = torch.tensor([[1, 2, 3]], dtype=torch.long, device=model.device) + cache_position = torch.arange(input_ids.shape[-1], dtype=torch.long, device=model.device) + + # Export + exported = exportable_module.export( + input_ids=input_ids, + cache_position=cache_position + ) + ``` + + Export with inputs_embeds: + ```python + # Prepare embeddings + inputs_embeds = torch.randn(1, 3, 768, device=model.device) # batch_size=1, seq_len=3, hidden_size=768 + cache_position = torch.arange(inputs_embeds.shape[1], dtype=torch.long, device=model.device) + + # Export + exported = exportable_module.export( + inputs_embeds=inputs_embeds, + cache_position=cache_position + ) + ``` + """ + if not (input_ids is None) ^ (inputs_embeds is None): + raise ValueError("Need to specify either input_ids or inputs_embeds.") + + if hasattr(self.model, "base_model_prefix"): + base = getattr(self.model, self.model.base_model_prefix, self.model) + model_device = base.device + elif hasattr(self.model, "model"): + model_device = self.model.model.device + else: + model_device = "cpu" + logging.warning( + "TorchExportableModuleForDecoderOnlyLM.export Can't infer device from the model. Set to CPU by default." + ) + + if input_ids is not None: + input_kwargs = { + "input_ids": input_ids, + "cache_position": cache_position + if cache_position is not None + else torch.arange(input_ids.shape[-1], dtype=torch.long, device=model_device), + } + else: # inputs_embeds + input_kwargs = { + "inputs_embeds": inputs_embeds, + "cache_position": cache_position + if cache_position is not None + else torch.arange(inputs_embeds.shape[1], dtype=torch.long, device=model_device), + } + + exported_program = torch.export.export( + self.model, + args=(), + kwargs=input_kwargs, + dynamic_shapes=dynamic_shapes, + strict=strict if strict is not None else True, + ) + + return exported_program + + @staticmethod + def generate( + exported_program: torch.export.ExportedProgram, + tokenizer, + prompt: str, + max_new_tokens: int = 20, + do_sample: bool = False, + temperature: float = 1.0, + top_k: int = 50, + top_p: float = 1.0, + device: str = "cpu", + ) -> str: + """ + Generate a sequence of tokens using an exported program. + + Args: + exported_program (`torch.export.ExportedProgram`): The exported model being used for generate. + tokenizer: The tokenizer to use. + prompt (str): The input prompt. + max_new_tokens (int): Maximum number of new tokens to generate. + do_sample (bool): Whether to use sampling or greedy decoding. + temperature (float): The temperature for sampling. + top_k (int): The number of highest probability tokens to keep for top-k sampling. + top_p (float): The cumulative probability for nucleus sampling. + device (str): The device to use. + + Returns: + str: The generated text. + """ + # Get the module from the exported program + exported_module = exported_program.module() + + # Tokenize the prompt + input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device) + + # Initialize with the prompt + generated_ids = input_ids.clone() + + # Process the prompt tokens first + curr_position = 0 + for i in range(input_ids.shape[1]): + # Process one token at a time + curr_input_ids = input_ids[:, i : i + 1] + curr_cache_position = torch.tensor([curr_position], dtype=torch.long, device=device) + + # Forward pass + _ = exported_module(input_ids=curr_input_ids, cache_position=curr_cache_position) + curr_position += 1 + + # Generate new tokens + for _ in range(max_new_tokens): + # Get the last token as input + curr_input_ids = generated_ids[:, -1:] + curr_cache_position = torch.tensor([curr_position], dtype=torch.long, device=device) + + # Forward pass to get next token logits + outputs = exported_module(input_ids=curr_input_ids, cache_position=curr_cache_position) + + # Get the next token ID + if do_sample: + # Apply temperature + if temperature > 0: + logits = outputs / temperature + else: + logits = outputs + + # Apply top-k filtering + if top_k > 0: + indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] + logits[indices_to_remove] = float("-inf") + + # Apply top-p (nucleus) filtering + if top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) + + # Remove tokens with cumulative probability above the threshold + sorted_indices_to_remove = cumulative_probs > top_p + # Shift the indices to the right to keep also the first token above the threshold + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() + sorted_indices_to_remove[..., 0] = 0 + + # Scatter sorted tensors to original indexing + indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove) + logits[indices_to_remove] = float("-inf") + + # Sample from the filtered distribution + probs = torch.softmax(logits, dim=-1) + next_token_id = torch.multinomial(probs, num_samples=1) + else: + # Greedy decoding + next_token_id = outputs.argmax(dim=-1, keepdim=True) + + # Ensure next_token_id has the right shape before concatenation + if next_token_id.dim() > 2: + next_token_id = next_token_id.squeeze(-1) + + # Append to the generated sequence + generated_ids = torch.cat([generated_ids, next_token_id], dim=-1) + curr_position += 1 + + # Stop if we generate an EOS token + if next_token_id.item() == tokenizer.eos_token_id: + break + + # Decode the generated text + return tokenizer.decode(generated_ids[0], skip_special_tokens=True) + + +def get_head_shapes(config) -> tuple[int | list[int], int | list[int]]: + """Returns a tuple `(num_heads, head_dim)` containing either 2 ints, or a list of int with the value for each + layer.""" + # Gemma4 has different head_dim and num_heads depending on layer type + if hasattr(config, "global_head_dim"): + head_dim = [ + config.global_head_dim if layer == "full_attention" else config.head_dim + for layer in config.layer_types[: -config.num_kv_shared_layers] + ] + num_heads = [ + config.num_global_key_value_heads + if layer == "full_attention" and config.attention_k_eq_v + else config.num_key_value_heads + for layer in config.layer_types[: -config.num_kv_shared_layers] + ] + else: + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + num_heads = getattr(config, "num_key_value_heads", config.num_attention_heads) + + return num_heads, head_dim + + +class TorchExportableModuleWithStaticCache(torch.nn.Module): + """ + A recipe module designed to make a `PreTrainedModel` exportable with `torch.export`, + specifically for decoder-only LM to `StaticCache`. This module ensures that the + exported model is compatible with further lowering and execution in `ExecuTorch`. + + Note: + This class is specifically designed to support export process using `torch.export` + in a way that ensures the model can be further lowered and run efficiently in `ExecuTorch`. + """ + + def __init__( + self, + model: PreTrainedModel, + batch_size: int | None = None, + max_cache_len: int | None = None, + device: torch.device | None = None, + ) -> None: + """ + Initializes the wrapper module with the pretrained model. + + Args: + model (`PreTrainedModel`): The pretrained model to wrap. The model must have caching + enabled and use a 'static' caching implementation. + batch_size (`Optional[int]`): The batch size of the model. If not provided, we check if a value can be found + in `generation_config.cache_config` and otherwise we raise a ValueError. + max_cache_len (`Optional[int]`): The maximum cache length for generation. Same mechanism as `batch_size` if + not provided. + device (`Optional[torch.device]`): The device to use. If not provided, we check if a value can be found + in `generation_config.cache_config` and otherwise we use `model.device` (no error is raised). + + Raises: + AssertionError: If the pretrained model does not have caching enabled or if it does + not use a 'static' caching implementation in `model.generation_config`. + ValueError: If `batch_size` or `max_cache_len` is not provided, either as an argument or in `cache_config`. + """ + super().__init__() + + config = model.config.get_text_config() + generation_config = model.generation_config + + # Sanity checks + if generation_config is None: + raise AssertionError( + "The model must have a generation config to be exported with static caching. " + "Please set `generation_config` in `model`." + ) + if not generation_config.use_cache: + raise AssertionError( + "The model must have caching enabled to be exported with static caching. " + "Please set `generation_config.use_cache=True`." + ) + if generation_config.cache_implementation != "static": + raise AssertionError( + "The model must use a 'static' caching implementation to be exported with static caching. " + "Please set `generation_config.cache_implementation='static'`." + ) + + cache_config = {} if generation_config.cache_config is None else generation_config.cache_config + + # Ensure batch_size and max_cache_len are set + if batch_size is None: + batch_size = cache_config.get("batch_size", None) + if batch_size is None: + raise ValueError("batch_size must be provided, either as an argument or in cache_config.") + if max_cache_len is None: + max_cache_len = cache_config.get("max_cache_len", None) + if max_cache_len is None: + raise ValueError("max_cache_len must be provided, either as an argument or in cache_config.") + # Infer device if not provided + if device is None: + device = cache_config.get("device", model.device) + + # Initialize the static cache + self.model = model + self.static_cache = StaticCache(max_cache_len=max_cache_len, config=config) + # Since StaticSlidingWindow have dynamic control flow that cannot be avoided, we have to replace them here by + # simple StaticLayer... It means that any generation beyond the window is unfortunately unsupported + for i, layer in enumerate(self.static_cache.layers): + if isinstance(layer, StaticSlidingWindowLayer): + self.static_cache.layers[i] = StaticLayer(max_cache_len) + num_heads, head_dim = get_head_shapes(config) + dtype = self.model.dtype + # We need this call to initialize all the layers (otherwise it's done lazily, which is not exportable) + self.static_cache.early_initialization(batch_size, num_heads, head_dim, dtype, device) + + # Register cache buffers to make them exportable + for i, layer in enumerate(self.static_cache.layers): + self.register_buffer(f"key_cache_{i}", layer.keys, persistent=False) + self.register_buffer(f"value_cache_{i}", layer.values, persistent=False) + self.register_buffer(f"cumulative_length_{i}", layer.cumulative_length, persistent=False) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, + cache_position: torch.Tensor | None = None, + ): + """ + Forward pass of the module, which is compatible with the ExecuTorch runtime. + + Args: + input_ids (`torch.Tensor`): Tensor representing current input token id to the module. + inputs_embeds (`torch.Tensor`): Tensor representing current input embeddings to the module. + cache_position (`torch.Tensor`): Tensor representing current input position in the cache. + + Returns: + torch.Tensor: Logits output from the model. + + This forward adapter serves two primary purposes: + + 1. **Making the Model `torch.export`-Compatible**: + The adapter hides unsupported objects, such as the `Cache`, from the graph inputs and outputs, + enabling the model to be exportable using `torch.export` without encountering issues. + + 2. **Ensuring Compatibility with `ExecuTorch` runtime**: + The adapter matches the model's forward signature with that in `executorch/extension/llm/runner`, + ensuring that the exported model can be executed in `ExecuTorch` out-of-the-box. + """ + # Start by resetting static cache (it's needed to be able to run several generations with the same exported program, + # as otherwise it's mutated in-place indefinitely - we cannot call reset in-between the `generate` as the program was + # already exported) + for layer in self.static_cache.layers: + layer.cumulative_length.copy_(cache_position[0:1]) + + past_key_values = self.static_cache + + outs = self.model( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + attention_mask=None, + past_key_values=past_key_values, + use_cache=True, + ) + if hasattr(outs, "logits"): + # Returned outputs is `CausalLMOutputWithPast` + return outs.logits + else: + # Returned the `last_hidden_state` from `BaseModelOutputWithPast` + return outs.last_hidden_state + + @staticmethod + def generate( + exported_program: torch.export.ExportedProgram, + prompt_token_ids: torch.Tensor, + max_new_tokens: int, + ) -> torch.Tensor: + """ + Generate a sequence of tokens using an exported program. + + This util function is designed to test exported models by simulating the generation process. + It processes the input prompt tokens sequentially (no parallel prefill). + This generate function is not intended to replace the original `generate` method, and the support + for leveraging the original `generate` is potentially planned! + + Args: + exported_program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`. + prompt_token_ids (`torch.Tensor`): Tensor representing the input prompt token IDs. + max_new_tokens (`int`): Maximum number of new tokens to generate. Note that the total generation + length is limited by both `max_new_tokens` and the model's cache size. + + Returns: + torch.Tensor: A tensor containing the generated sequence of token IDs, including the original prompt tokens. + """ + device = prompt_token_ids.device + prompt_token_len = prompt_token_ids.shape[-1] + max_generation_length = prompt_token_len + max_new_tokens + for buffer_name, buffer in exported_program.named_buffers(): + if buffer_name.startswith("key_cache"): + max_cache_len = buffer.shape[2] + max_generation_length = min(max_generation_length, max_cache_len) + break + + response_tokens = [] + for input_pos in range(min(max_generation_length, prompt_token_len)): + result = exported_program.module().forward( + input_ids=prompt_token_ids[:, input_pos : input_pos + 1], + cache_position=torch.tensor([input_pos], dtype=torch.long, device=device), + ) + response_tokens.append(prompt_token_ids[0][input_pos].item()) + + current_token = torch.argmax(result[:, -1, :], dim=-1).item() + response_tokens.append(current_token) + + while len(response_tokens) < max_generation_length: + result = exported_program.module().forward( + input_ids=torch.tensor([[current_token]], dtype=torch.long, device=device), + cache_position=torch.tensor([len(response_tokens)], dtype=torch.long, device=device), + ) + current_token = torch.argmax(result[:, -1, :], dim=-1).item() + response_tokens.append(current_token) + + return torch.tensor([response_tokens], dtype=torch.long, device=device) + + +class TorchExportableModuleWithHybridCache(torch.nn.Module): + """ + A recipe module designed to make a `PreTrainedModel` exportable with `torch.export`, + specifically for decoder-only LM to hybrid `StaticCache`. This module ensures that the + exported model is compatible with further lowering and execution in `ExecuTorch`. + """ + + def __init__( + self, + model: PreTrainedModel, + batch_size: int | None = None, + max_cache_len: int | None = None, + device: torch.device | None = None, + ) -> None: + """ + Initializes the exportable module. + + Args: + model (`PreTrainedModel`): The pretrained model to wrap. + batch_size (`Optional[int]`): The batch size of the model. If not provided, we check if a value can be found + in `generation_config.cache_config` and otherwise we raise a ValueError. + max_cache_len (`Optional[int]`): The maximum cache length for generation. Same mechanism as `batch_size` if + not provided. + device (`Optional[torch.device]`): The device to use. If not provided, we check if a value can be found + in `generation_config.cache_config` and otherwise we use `model.device` (no error is raised). + Raises: + AssertionError: If the model doesn't have the expected configuration for hybrid StaticCache. + ValueError: If `batch_size` or `max_cache_len` is not provided, either as an argument or in `cache_config`. + """ + super().__init__() + self.model = model + config = model.config.get_text_config() + generation_config = model.generation_config + + # Sanity checks + if generation_config is None: + raise AssertionError( + "The model must have a generation config to be exported with static caching. " + "Please set `generation_config` in `model`." + ) + if not config.use_cache: + raise AssertionError("Model must have caching enabled.") + + cache_config = {} if generation_config.cache_config is None else generation_config.cache_config + # Ensure batch_size and max_cache_len are set + if batch_size is None: + batch_size = cache_config.get("batch_size", None) + if batch_size is None: + raise ValueError("batch_size must be provided, either as an argument or in cache_config.") + if max_cache_len is None: + max_cache_len = cache_config.get("max_cache_len", None) + if max_cache_len is None: + raise ValueError("max_cache_len must be provided, either as an argument or in cache_config.") + # Infer device if not provided + if device is None: + device = cache_config.get("device", model.device) + + # Initialize the cache + self.cache = StaticCache(config=config, max_cache_len=max_cache_len) + # Since StaticSlidingWindow have dynamic control flow that cannot be avoided, we have to replace them here by + # simple StaticLayer... It means that any generation beyond the window is unfortunately unsupported + for i, layer in enumerate(self.cache.layers): + if isinstance(layer, StaticSlidingWindowLayer): + self.cache.layers[i] = StaticLayer(max_cache_len) + num_heads, head_dim = get_head_shapes(config) + dtype = self.model.dtype + # We need this call to initialize all the layers (otherwise it's done lazily, which is not exportable) + self.cache.early_initialization(batch_size, num_heads, head_dim, dtype, device) + + # Register cache buffers to make them exportable + for i, layer in enumerate(self.cache.layers): + self.register_buffer(f"key_cache_{i}", layer.keys, persistent=False) + self.register_buffer(f"value_cache_{i}", layer.values, persistent=False) + self.register_buffer(f"cumulative_length_{i}", layer.cumulative_length, persistent=False) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, + cache_position: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + Forward pass of the module, which is compatible with the ExecuTorch llm runner. + + Args: + input_ids (`torch.Tensor`): Tensor representing current input token id to the module. + inputs_embeds (`Optional[torch.Tensor]`): Tensor representing current input embeddings to the module. + cache_position (`torch.Tensor`): Tensor representing current input position in the cache. + + Returns: + torch.Tensor: Logits output from the model. + """ + # Start by resetting static cache (it's needed to be able to run several generations with the same exported program, + # as otherwise it's mutated in-place indefinitely - we cannot call reset in-between the `generate` as the program was + # already exported) + for layer in self.cache.layers: + layer.cumulative_length.copy_(cache_position[0:1]) + + # Forward pass with the model + outputs = self.model( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + attention_mask=None, + past_key_values=self.cache, + use_cache=True, + ) + + # Return only the logits to simplify the export + return outputs.logits + + +def convert_and_export_with_cache( + model: PreTrainedModel, + example_input_ids: torch.Tensor | None = None, + example_cache_position: torch.Tensor | None = None, + dynamic_shapes: dict | None = None, + strict: bool | None = None, +): + """ + Convert a `PreTrainedModel` into an exportable module and export it using `torch.export`, + ensuring the exported model is compatible with `ExecuTorch`. + + Args: + model (`PreTrainedModel`): The pretrained model to be exported. + example_input_ids (`Optional[torch.Tensor]`): Example input token id used by `torch.export`. + example_cache_position (`Optional[torch.Tensor]`): Example current cache position used by `torch.export`. + dynamic_shapes(`Optional[dict]`): Dynamic shapes used by `torch.export`. + strict(`Optional[bool]`): Flag to instruct `torch.export` to use `torchdynamo`. + + Returns: + Exported program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`. + """ + + import torch.export._trace + + with torch.no_grad(): + # TODO: The default inputs only work for text models. We need to add support for vision/audio models. + example_input_ids = ( + example_input_ids + if example_input_ids is not None + else torch.tensor([[1]], dtype=torch.long, device=model.device) + ) + example_cache_position = ( + example_cache_position + if example_cache_position is not None + else torch.tensor([0], dtype=torch.long, device=model.device) + ) + + if is_torch_greater_or_equal("2.6.0"): + exported_program = torch.export.export( + TorchExportableModuleWithStaticCache(model), + args=(), + kwargs={"input_ids": example_input_ids, "cache_position": example_cache_position}, + dynamic_shapes=dynamic_shapes, + strict=strict if strict is not None else True, + ) + else: + if dynamic_shapes is not None: + logging.warning( + "Dynamic shapes spec will be ignored by convert_and_export_with_cache for torch < 2.6.0." + ) + if strict is not None: + logging.warning("The strict flag will be ignored by convert_and_export_with_cache for torch < 2.6.0.") + # We have to keep this path for BC. + # + # Due to issue https://github.com/pytorch/pytorch/issues/128394, we need to switch to use an internal + # export API and pre_dispatch=False. Switch to use the public API once the issue is included in 2.5 release. + exported_program = torch.export._trace._export( + TorchExportableModuleWithStaticCache(model), + args=(), + kwargs={"input_ids": example_input_ids, "cache_position": example_cache_position}, + pre_dispatch=False, + strict=True, + ) + return exported_program + + +class Seq2SeqLMEncoderExportableModule(torch.nn.Module): + """ + A wrapper module designed to make a Seq2Seq LM encoder exportable with `torch.export`. + This module ensures that the exported encoder model is compatible with ExecuTorch. + """ + + def __init__(self, encoder_model): + super().__init__() + self.encoder = encoder_model + + def forward(self, input_ids): + return self.encoder(input_ids=input_ids).last_hidden_state + + +class Seq2SeqLMDecoderExportableModuleWithStaticCache(torch.nn.Module): + """ + A wrapper module designed to make a Seq2Seq LM decoder exportable with `torch.export`, + specifically for use with static caching. This module ensures the exported decoder + is compatible with ExecuTorch. + """ + + def __init__(self, model, max_static_cache_length, batch_size): + super().__init__() + + # Get the decoder component + self.decoder = model.get_decoder() + self.lm_head = model.lm_head + self.config = model.config + + # Detect the device of the exported models by checking a parameter + # We'll use the model's device as the target device + model_device = next(model.parameters()).device + + # Initialize static cache for decoder and DynamicCache for encoder + self.static_cache = StaticCache(config=self.config, max_cache_len=max_static_cache_length) + # Since StaticSlidingWindow have dynamic control flow that cannot be avoided, we have to replace them here by + # simple StaticLayer... It means that any generation beyond the window is unfortunately unsupported + for i, layer in enumerate(self.static_cache.layers): + if isinstance(layer, StaticSlidingWindowLayer): + self.static_cache.layers[i] = StaticLayer(max_static_cache_length) + num_heads, head_dim = get_head_shapes(self.config) + self.static_cache.early_initialization(batch_size, num_heads, head_dim, torch.float32, model_device) + self.cache = EncoderDecoderCache(self.static_cache, DynamicCache(config=self.config)) + + register_dynamic_cache_export_support() + + # Register cache buffers to make them exportable + for i, layer in enumerate(self.static_cache.layers): + self.register_buffer(f"key_cache_{i}", layer.keys, persistent=False) + self.register_buffer(f"value_cache_{i}", layer.values, persistent=False) + self.register_buffer(f"cumulative_length_{i}", layer.cumulative_length, persistent=False) + + def forward(self, decoder_input_ids, encoder_hidden_states, cache_position): + # Start by resetting static cache (it's needed to be able to run several generations with the same exported program, + # as otherwise it's mutated in-place indefinitely - we cannot call reset in-between the `generate` as the program was + # already exported) + for layer in self.static_cache.layers: + layer.cumulative_length.copy_(cache_position[0:1]) + + # Get outputs from decoder + outputs = self.decoder( + input_ids=decoder_input_ids, + encoder_hidden_states=encoder_hidden_states, + past_key_values=self.cache, + use_cache=True, + ) + + # Apply language model head + lm_logits = self.lm_head(outputs[0]) + + return lm_logits + + +class Seq2SeqLMExportableModule(torch.nn.Module): + def __init__( + self, model, batch_size=1, max_hidden_seq_length=4096, cache_implementation="static", max_cache_length=1024 + ): + super().__init__() + + self.full_model = model + self.encoder = model.get_encoder() + self.config = model.config + self.max_hidden_seq_length = max_hidden_seq_length + self.generation_config = GenerationConfig( + use_cache=True, + max_length=max_cache_length, + cache_implementation=cache_implementation, + cache_config={ + "batch_size": batch_size, + "max_cache_len": max_cache_length, + }, + eos_token_id=model.generation_config.eos_token_id, + ) + self.exported_encoder = None + self.exported_decoder = None + + def _export_encoder(self, encoder_input_ids): + wrapped_encoder = Seq2SeqLMEncoderExportableModule(self.encoder).to(self.full_model.device).eval() + + # Define dynamic sequence length for encoder + seq_len_dim = torch.export.Dim("encoder_seq_length", max=self.max_hidden_seq_length) + + # Export the encoder + with torch.no_grad(): + exported_encoder = torch.export.export( + wrapped_encoder, (encoder_input_ids,), dynamic_shapes={"input_ids": {1: seq_len_dim}}, strict=True + ) + + return exported_encoder + + def _export_decoder(self, decoder_input_ids, encoder_hidden_states, cache_position): + target_device = self.full_model.device + wrapped_decoder = ( + Seq2SeqLMDecoderExportableModuleWithStaticCache( + model=self.full_model, + max_static_cache_length=self.generation_config.cache_config.get("max_cache_len"), + batch_size=self.generation_config.cache_config.get("batch_size"), + ) + .to(target_device) + .eval() + ) + + # Move input tensors to the same device as the wrapped decoder + decoder_input_ids = decoder_input_ids.to(target_device) + encoder_hidden_states = encoder_hidden_states.to(target_device) + cache_position = cache_position.to(target_device) + + # Define dynamic dimension for encoder output sequence length + encoder_seq_len_dim = torch.export.Dim("encoder_hidden_seq_length", max=self.max_hidden_seq_length) + + # Export the decoder + with torch.no_grad(): + exported_decoder = torch.export.export( + wrapped_decoder, + (decoder_input_ids, encoder_hidden_states, cache_position), + dynamic_shapes={ + "decoder_input_ids": None, + "encoder_hidden_states": {1: encoder_seq_len_dim}, + "cache_position": None, + }, + strict=True, + ) + + return exported_decoder + + def export(self, encoder_input_ids=None, decoder_input_ids=None, encoder_hidden_states=None, cache_position=None): + device = self.full_model.device + example_encoder_input_ids = ( + encoder_input_ids + if encoder_input_ids is not None + else torch.ones((1, 10), dtype=torch.long, device=device) + ) + example_decoder_input_ids = ( + decoder_input_ids + if decoder_input_ids is not None + else torch.tensor([[0]], dtype=torch.long, device=device) + ) # Start token + example_cache_position = ( + cache_position if cache_position is not None else torch.tensor([0], dtype=torch.long, device=device) + ) + example_encoder_hidden_states = ( + encoder_hidden_states + if encoder_hidden_states is not None + else torch.zeros( + (self.generation_config.cache_config.get("batch_size"), 10, self.config.d_model), + dtype=torch.float32, + device=device, + ) + ) + self.exported_encoder = self._export_encoder(example_encoder_input_ids) + self.exported_decoder = self._export_decoder( + example_decoder_input_ids, example_encoder_hidden_states, example_cache_position + ) + + # Return self to allow chaining + return self + + def generate(self, prompt_token_ids, max_new_tokens): + with torch.no_grad(): + model_device = self.full_model.device + + # Move input to the model's device if it's on a different device + if prompt_token_ids.device != model_device: + prompt_token_ids = prompt_token_ids.to(model_device) + + # Run encoder + encoder_output = self.exported_encoder.module()(prompt_token_ids) + + # Initialize with start token (0 for T5) on the correct device + decoder_input_ids = torch.tensor([[0]], dtype=torch.long, device=model_device) + generated_ids = [0] + + # Generate tokens one by one + for i in range(max_new_tokens - 1): + # Run decoder for next token prediction + logits = self.exported_decoder.module()( + decoder_input_ids, encoder_output, torch.tensor([i], dtype=torch.long, device=model_device) + ) + + # Get next token + next_token = torch.argmax(logits[:, -1, :], dim=-1).item() + generated_ids.append(next_token) + + # Update input for next iteration on the correct device + decoder_input_ids = torch.tensor([[next_token]], dtype=torch.long, device=model_device) + + # Check if EOS token + if next_token == self.generation_config.eos_token_id: + break + + return generated_ids + + +def export_with_dynamic_cache( + model: PreTrainedModel, + example_input_ids: torch.Tensor | None = None, + example_attention_mask: torch.Tensor | None = None, +): + """ + Export a model with DynamicCache using `torch.export`, ensuring the exported model is compatible with `ExecuTorch`. + + Args: + model (`PreTrainedModel`): The pretrained model to be exported. + example_input_ids (`Optional[torch.Tensor]`): Example input token id used by `torch.export`. + example_attention_mask (`Optional[torch.Tensor]`): Example attention mask used by `torch.export`. + + Returns: + Exported program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`. + """ + + register_dynamic_cache_export_support() + + with torch.no_grad(): + exported_program = torch.export.export( + model, + (), + { + "input_ids": example_input_ids, + "attention_mask": example_attention_mask, + "past_key_values": DynamicCache(config=model.config), + "use_cache": True, + }, + strict=False, + ) + return exported_program + + +def register_dynamic_cache_export_support(): + """ + Utilities for `DynamicCache` <> torch.export support + """ + + try: + torch.utils._pytree.register_pytree_node( + DynamicCache, + lambda dynamic_cache: torch.utils._pytree._dict_flatten(_get_cache_dict(dynamic_cache)), + _unflatten_dynamic_cache, + serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}", + flatten_with_keys_fn=lambda dynamic_cache: torch.utils._pytree._dict_flatten_with_keys( + _get_cache_dict(dynamic_cache) + ), + ) + # TODO (tmanlaibaatar) This won't be needed in torch 2.7. + torch.fx._pytree.register_pytree_flatten_spec( + DynamicCache, + lambda cache, spec: torch.fx._pytree._dict_flatten_spec(_get_cache_dict(cache), spec), + ) + # Catching this in case there are multiple runs for some test runs + except ValueError as e: + if "already registered as pytree node" not in str(e): + raise + + +def _get_cache_dict(cache: DynamicCache): + """Convert cache to dictionary format for pytree operations.""" + if any(not isinstance(layer, (DynamicLayer, DynamicSlidingWindowLayer)) for layer in cache.layers): + raise RuntimeError("This pytree flattening function should only be applied to DynamicCache") + + if not is_torch_greater_or_equal_than_2_6: + logging.warning("DynamicCache + torch.export is tested on torch 2.6.0+ and may not work on earlier versions.") + + return { + "key_cache": [layer.keys for layer in cache.layers if layer.keys is not None], + "value_cache": [layer.values for layer in cache.layers if layer.values is not None], + } + + +def _unflatten_dynamic_cache(values, context: torch.utils._pytree.Context): + dictionary = torch.utils._pytree._dict_unflatten(values, context) + cache = DynamicCache() + # Reconstruct layers from keys and values lists + key_list = dictionary.get("key_cache", []) + value_list = dictionary.get("value_cache", []) + for idx in range(max(len(key_list), len(value_list))): + key = key_list[idx] if idx < len(key_list) else None + value = value_list[idx] if idx < len(value_list) else None + cache.update(key, value, idx) + return cache diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/hub_kernels.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/hub_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..76344e5453e57c55d019fa13806b0b2997199057 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/hub_kernels.py @@ -0,0 +1,514 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import importlib.metadata +import os +import re +from collections.abc import Callable +from contextlib import contextmanager +from types import ModuleType + +from packaging import version as pkg_version + +from ..utils import ENV_VARS_TRUE_VALUES, logging +from ..utils.import_utils import is_kernels_available +from .flash_attention import flash_attention_forward + + +logger = logging.get_logger(__name__) + +try: + from kernels import ( + Device, + LayerRepository, + Mode, + register_kernel_mapping, + replace_kernel_forward_from_hub, + ) + from kernels import ( + get_kernel as get_kernel_hub, + ) + from kernels import ( + use_kernel_forward_from_hub as _kernels_use_kernel_forward_from_hub, + ) + + # Try to import FuncRepository, fallback if not available + try: + from kernels import FuncRepository + except ImportError: + FuncRepository = None + + # Try to import use_kernel_func_from_hub, fallback if not available + try: + from kernels import use_kernel_func_from_hub as _kernels_use_kernel_func_from_hub + + _has_use_kernel_func_from_hub = True + except ImportError: + _has_use_kernel_func_from_hub = False + + _TRANSFORMERS_USE_HUB_KERNELS = os.environ.get("USE_HUB_KERNELS", "YES").upper() + _kernels_available = True + _kernels_enabled = _TRANSFORMERS_USE_HUB_KERNELS in ENV_VARS_TRUE_VALUES + + def use_kernel_forward_from_hub(layer_name: str): + if _kernels_enabled: + return _kernels_use_kernel_forward_from_hub(layer_name) + else: + logger.warning_once( + f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}" + ) + return lambda cls: cls + + def use_kernel_func_from_hub(func_name: str): + if _kernels_enabled and _has_use_kernel_func_from_hub: + return _kernels_use_kernel_func_from_hub(func_name) + else: + if not _has_use_kernel_func_from_hub: + logger.warning_once( + "use_kernel_func_from_hub is not available in the installed kernels version. " + "Please upgrade kernels to use this feature." + ) + else: + logger.warning_once( + f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}" + ) + return lambda func: func + + _KERNEL_MAPPING: dict[str, dict[Device | str, LayerRepository | dict[Mode, LayerRepository]]] = { + "MultiScaleDeformableAttention": { + "cuda": LayerRepository( + repo_id="kernels-community/deformable-detr", + layer_name="MultiScaleDeformableAttention", + ) + }, + "Llama4TextMoe": { + "cuda": LayerRepository( + repo_id="kernels-community/moe", + layer_name="Llama4TextMoe", + ) + }, + "RMSNorm": { + "cuda": { + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/liger_kernels", + layer_name="LigerRMSNorm", + # revision="pure-layer-test", + ), + }, + "rocm": { + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/liger_kernels", + layer_name="LigerRMSNorm", + ) + }, + "xpu": { + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/rmsnorm", + layer_name="RMSNorm", + ) + }, + "mps": { + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/mlx_rmsnorm", + layer_name="RMSNorm", + ) + }, + "npu": { + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/liger_kernels", + layer_name="LigerRMSNorm", + ) + }, + }, + "MLP": { + "cuda": LayerRepository( + repo_id="medmekk/triton-llama-mlp", + layer_name="TritonLlamaMLP", + ) + }, + "MegaBlocksMoeMLP": { + "cuda": { + Mode.TRAINING: LayerRepository( + repo_id="kernels-community/megablocks", + layer_name="MegaBlocksMoeMLP", + ), + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/megablocks", + layer_name="MegaBlocksMoeMLP", + ), + }, + "rocm": { + Mode.INFERENCE: LayerRepository( + repo_id="ahadnagy/megablocks", + layer_name="MegaBlocksMoeMLP", + ) + }, + "xpu": { + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/megablocks", + layer_name="MegaBlocksMoeMLP", + ) + }, + "cpu": { + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/megablocks", + layer_name="CPUMegaBlocksMoeMLP", + ) + }, + }, + "FastGELU": { + "cuda": { + Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/activation", + layer_name="FastGELU", + version=1, + ) + } + }, + "QuickGELU": { + "cuda": { + Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/activation", + layer_name="QuickGELU", + version=1, + ) + } + }, + "NewGELU": { + "cuda": { + Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/activation", + layer_name="NewGELU", + version=1, + ) + } + }, + "SiLU": { + "cuda": { + Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/activation", layer_name="Silu", version=1 + ) + } + }, + "GeLU": { + "cuda": { + Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/activation", layer_name="Gelu", version=1 + ) + } + }, + "GeluTanh": { + "cuda": { + Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/activation", layer_name="GeluTanh", version=1 + ) + } + }, + } + + # Add function kernel mappings if FuncRepository is available + if FuncRepository is not None: + _KERNEL_MAPPING["rotary_pos_emb"] = { + "xpu": { + Mode.INFERENCE: FuncRepository( + repo_id="kernels-community/rotary", func_name="apply_rotary_transformers" + ) + }, + "cuda": { + Mode.TRAINING: FuncRepository( + repo_id="kernels-community/rotary", func_name="apply_rotary_transformers" + ), + Mode.INFERENCE: FuncRepository( + repo_id="kernels-community/rotary", func_name="apply_rotary_transformers" + ), + }, + } + + def has_key(d, key): + return key in d or any(isinstance(v, dict) and has_key(v, key) for v in d.values()) + + def register_kernel_mapping_transformers(mapping=None): + if mapping is None: + mapping = _KERNEL_MAPPING + if has_key(mapping, "xpu") and not is_kernels_available(MIN_VERSION="0.10.2"): + raise ImportError( + "kernels uses an incompatible version. Please install the latest version with `pip install -U kernels`." + ) + register_kernel_mapping(mapping) + + +except ImportError: + _kernels_available = False + _kernels_enabled = False + + # Stub to make decorators int transformers work when `kernels` + # is not installed. + def use_kernel_forward_from_hub(*args, **kwargs): + def decorator(cls): + return cls + + return decorator + + def use_kernel_func_from_hub(*args, **kwargs): + def decorator(func): + return func + + return decorator + + class LayerRepository: + def __init__(self, *args, **kwargs): + raise RuntimeError("LayerRepository requires `kernels` to be installed. Run `pip install kernels`.") + + def replace_kernel_forward_from_hub(*args, **kwargs): + raise RuntimeError( + "replace_kernel_forward_from_hub requires `kernels` to be installed. Run `pip install kernels`." + ) + + def register_kernel_mapping(*args, **kwargs): + raise RuntimeError("register_kernel_mapping requires `kernels` to be installed. Run `pip install kernels`.") + + def register_kernel_mapping_transformers(*args, **kwargs): + raise RuntimeError( + "register_kernel_mapping_transformers requires `kernels` to be installed. Run `pip install kernels`." + ) + + +_HUB_KERNEL_MAPPING: dict[str, dict[str, str]] = { + "causal-conv1d": {"repo_id": "kernels-community/causal-conv1d", "version": 1}, + "mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1}, + "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1}, + "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 1}, + "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 1}, + "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"}, +} + +_KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {} + + +def is_kernel(attn_implementation: str | None) -> bool: + """Check whether `attn_implementation` matches a kernel pattern from the hub.""" + return ( + attn_implementation is not None + and re.search(r"^[^/:]+/[^/:]+(?:@[^/:]+)?(?::[^/:]+)?$", attn_implementation) is not None + ) + + +def load_and_register_attn_kernel( + attn_implementation: str, attention_wrapper: Callable | None = None, allow_all_kernels: bool = False +) -> ModuleType | None: + """ + Load and register the kernel associated to `attn_implementation`. + + Args: + attn_implementation: A string, usually a kernel repo like "kernels-community/flash-mla". + attn_wrapper: a callable for the wrapper around the attention implementation. In `transformers` we + have a wrapper around the `flash_attn_var_len` call, and the same goes for `sdpa` and `eager`. + They just prepare the arguments properly. This is mostly used for continious batching, where we + want the `paged` wrapper, which calls the paged cache. + allow_all_kernels (`bool`, optional): + Whether to load kernels from unverified hub repos, if it is a custom kernel outside of the `kernels-community` + hub repository. + """ + from ..masking_utils import ALL_MASK_ATTENTION_FUNCTIONS + from ..modeling_utils import ALL_ATTENTION_FUNCTIONS + + actual_attn_name = attn_implementation.split("|")[1] if "|" in attn_implementation else attn_implementation + if not is_kernel(actual_attn_name): + return None + if not _kernels_available: + raise ImportError( + "`kernels` is either not installed or uses an incompatible version. " + "Please install the latest version with `pip install -U kernels`." + ) + + # Extract repo_id and kernel_name from the string + if ":" in actual_attn_name: + repo_id, kernel_name = actual_attn_name.split(":") + kernel_name = kernel_name.strip() + else: + repo_id = actual_attn_name + kernel_name = None + repo_id = repo_id.strip() + # extract the rev after the @ if it exists + repo_id, _, rev = repo_id.partition("@") + repo_id = repo_id.strip() + rev = rev.strip() if rev else None + + # Load the kernel from hub + try: + kernel = get_kernel(repo_id, revision=rev, allow_all_kernels=allow_all_kernels) + except Exception as e: + raise ValueError(f"An error occurred while trying to load from '{repo_id}': {e}.") + # correctly wrap the kernel + if hasattr(kernel, "flash_attn_varlen_func"): + if attention_wrapper is None: + attention_wrapper = flash_attention_forward + kernel_function = attention_wrapper + elif kernel_name is not None: + kernel_function = getattr(kernel, kernel_name) + + # Register the kernel as a valid attention + ALL_ATTENTION_FUNCTIONS.register(attn_implementation, kernel_function) + ALL_MASK_ATTENTION_FUNCTIONS.register(attn_implementation, ALL_MASK_ATTENTION_FUNCTIONS["flash_attention_2"]) + + return kernel + + +def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _KERNEL_MODULE_MAPPING): + if kernel_name in mapping and isinstance(mapping[kernel_name], ModuleType): + return mapping[kernel_name] + if kernel_name not in _HUB_KERNEL_MAPPING: + logger.warning_once(f"Kernel {kernel_name} not found in _HUB_KERNEL_MAPPING") + mapping[kernel_name] = None + return None + if _kernels_available: + try: + repo_id = _HUB_KERNEL_MAPPING[kernel_name]["repo_id"] + revision = _HUB_KERNEL_MAPPING[kernel_name].get("revision", None) + version = _HUB_KERNEL_MAPPING[kernel_name].get("version", None) + kernel = get_kernel(repo_id, revision=revision, version=version) + mapping[kernel_name] = kernel + except FileNotFoundError as e: + mapping[kernel_name] = None + logger.warning_once(f"Failed to load kernel {kernel_name}: {e}") + except AssertionError: + # Happens when torch is built without an accelerator backend; fall back to slow path. + mapping[kernel_name] = None + + else: + # Try to import is_{kernel_name}_available from ..utils + import importlib + + new_kernel_name = kernel_name.replace("-", "_") + func_name = f"is_{new_kernel_name}_available" + + try: + utils_mod = importlib.import_module("..utils.import_utils", __package__) + is_kernel_available = getattr(utils_mod, func_name, None) + except Exception: + is_kernel_available = None + + if callable(is_kernel_available) and is_kernel_available(): + # Try to import the module "{kernel_name}" from parent package level + try: + module = importlib.import_module(f"{new_kernel_name}") + mapping[kernel_name] = module + return module + except Exception: + mapping[kernel_name] = None + else: + mapping[kernel_name] = None + + return mapping[kernel_name] + + +def get_kernel( + kernel_name: str, + revision: str | None = None, + version: int | str | None = None, + allow_all_kernels: bool = False, +) -> ModuleType: + from .. import __version__ + + if not _kernels_available: + raise ImportError( + "`kernels` is either not installed or uses an incompatible version. Please install the latest version " + "with `pip install -U kernels`." + ) + + repo_parent = kernel_name.split("/")[0] + # all `kernels-community` repos are trusted by default! + if repo_parent != "kernels-community" and not allow_all_kernels: + raise ValueError( + "You need to specify `allow_all_kernels=True` to use kernels outside of the `kernels-community` repository" + ) + + user_agent = {"framework": "transformers", "version": __version__, "repo_id": kernel_name} + kernels_version = importlib.metadata.version("kernels") + if pkg_version.parse(kernels_version) >= pkg_version.parse("0.10.4"): + return get_kernel_hub(kernel_name, revision=revision, version=version, user_agent=user_agent) + else: + return get_kernel_hub(kernel_name, revision=revision, version=version) + + +def use_kernelized_func(module_names: list[Callable] | Callable): + """ + This decorator attaches the target function within the module as a plain attribute (not as a submodule). + Keep in mind that this registration is only meant for `kernelize` to recognize its target modules (i.e. + function exchanged for a weightless `nn.Module` with the same forward) to then exchange to the kernel + variation (in-place) if the conditions are met. + + We cache each of these function-based registrations: After proper registration and exchange it is removed + from the module's `_modules` dict as it does not really act as `nn.Module` but a base function. + """ + if isinstance(module_names, Callable): + module_names = [module_names] + + def decorator(cls): + orig_init = cls.__init__ + + def new_init(self, *args, **kwargs): + orig_init(self, *args, **kwargs) + + # Register new function as non-submodule within the modules dict + hidden_kernels = self.__dict__.setdefault("_hidden_kernels", {}) + for fn in module_names: + name = ( + getattr(fn, "__name__", None) + or getattr(fn, "kernel_layer_name", None) + or getattr(fn, "func_name", None) + ) + if name is None: + raise ValueError(f"Could not infer kernel function name for {fn!r}") + + # Do not register as submodule! Hide it behind a dict to be removed later after registering it + hidden_kernels[name] = fn + + cls.__init__ = new_init + return cls + + return decorator + + +# Whether to allow hub kernels coming from untrusted repos, i.e. repos outside `kernels-community` +ALLOW_ALL_KERNELS = False + + +@contextmanager +def allow_all_hub_kernels(): + """ + Context manager used to adjust the value of the global `ALLOW_HUB_KERNELS`. This is needed, as this argument + cannot be forwarded directly to the `__init__` of the models, where we set the attention implementation. + """ + global ALLOW_ALL_KERNELS + + try: + ALLOW_ALL_KERNELS = True + + yield + finally: + # Set back the original + ALLOW_ALL_KERNELS = False + + +__all__ = [ + "LayerRepository", + "use_kernel_forward_from_hub", + "use_kernel_func_from_hub", + "register_kernel_mapping", + "register_kernel_mapping_transformers", + "replace_kernel_forward_from_hub", + "lazy_load_kernel", + "get_kernel", + "use_kernelized_func", +] # type: ignore diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/moe.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/moe.py new file mode 100644 index 0000000000000000000000000000000000000000..4f107f03e6fb6f3695512dcbf1fe8da76364f398 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/moe.py @@ -0,0 +1,582 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from collections.abc import Callable +from functools import wraps + +from ..utils import logging +from ..utils.generic import GeneralInterface +from ..utils.import_utils import ( + is_torch_available, + is_torch_greater_or_equal, + is_torch_less_or_equal, + is_torchdynamo_compiling, +) +from .sonicmoe import sonicmoe_experts_forward + + +if is_torch_available(): + import torch + + # Patch the version-check helpers so dynamo doesn't trace into them — they transitively call + # `importlib.util.find_spec`, which dynamo refuses to trace. `assume_constant_result` makes + # dynamo evaluate them once at trace time and inline the bool, no body tracing. + is_torch_greater_or_equal = torch._dynamo.assume_constant_result(is_torch_greater_or_equal) + is_torch_less_or_equal = torch._dynamo.assume_constant_result(is_torch_less_or_equal) + + +logger = logging.get_logger(__name__) + + +# Examples of experts class with its eager mm implementation +# class Experts(torch.nn.Module): +# """Collection of expert weights stored as 3D tensors.""" + +# def __init__(self, config): +# super().__init__() +# self.num_experts = config.n_routed_experts +# self.hidden_dim = config.hidden_size +# self.intermediate_dim = config.moe_intermediate_size +# self.gate_up_proj = torch.nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) +# self.down_proj = torch.nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) +# self.act_fn = ACT2FN[config.hidden_act] + +# def forward( +# self, +# hidden_states: torch.Tensor, +# top_k_index: torch.Tensor, +# top_k_weights: torch.Tensor, +# ) -> torch.Tensor: +# final_hidden_states = torch.zeros_like(hidden_states) +# with torch.no_grad(): +# expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts) +# expert_mask = expert_mask.permute(2, 1, 0) +# expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + +# for expert_idx in expert_hit: +# expert_idx = expert_idx[0] +# if expert_idx == self.num_experts: +# continue +# top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) +# current_state = hidden_states[token_idx] +# gate, up = torch.nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1) +# current_hidden_states = self.act_fn(gate) * up +# current_hidden_states = torch.nn.functional.linear(current_hidden_states, self.down_proj[expert_idx]) +# current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None] +# final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) + +# return final_hidden_states + + +def _batched_linear( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + is_transposed: bool = False, +) -> torch.Tensor: + """Batched linear layer supporting optional bias and transposed weights. + + Args: + input (`torch.Tensor`): + Input tensor of shape (batch_size, input_dim). + weight (`torch.Tensor`): + Weight tensor of shape (batch_size, output_dim, input_dim) if transposed is `False`, + else of shape (batch_size, input_dim, output_dim). + bias (`torch.Tensor`, *optional*): + Bias tensor of shape (batch_size, output_dim). Default is `None`. + is_transposed (`bool`, *optional*, defaults to `False`): + Whether the weight tensor is transposed. + Returns: + `torch.Tensor`: Output tensor of shape (batch_size, output_dim). + """ + if is_transposed: + # (batch_size, 1, input_dim) @ (batch_size, input_dim, output_dim) -> (batch_size, 1, output_dim) -> (batch_size, output_dim) + out = torch.bmm(input.unsqueeze(1), weight).squeeze(1) + else: + # (batch_size, output_dim, input_dim) @ (batch_size, input_dim, 1) -> (batch_size, output_dim, 1) -> (batch_size, output_dim) + out = torch.bmm(weight, input.unsqueeze(-1)).squeeze(-1) + + if bias is not None: + out.add_(bias) + + return out + + +def batched_mm_experts_forward( + self: torch.nn.Module, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + num_top_k = top_k_index.size(-1) + num_tokens = hidden_states.size(0) + hidden_dim = hidden_states.size(-1) + + # S is the number of selected tokens-experts pairs (S = num_tokens * num_top_k) + # Replicate each token num_top_k times to align with the flattened (S,) routing tensors. + selected_hidden_states = hidden_states.repeat_interleave(num_top_k, dim=0) + sample_weights = top_k_weights.reshape(-1) # (S,) + expert_ids = top_k_index.reshape(-1) # (S,) + + # Clamp EP sentinels so `gate_up_proj[expert_ids]` stays in-bounds. Routing weights are already + # zero at sentinel slots (RouterParallel masks them at dispatch), so the weighted mul drops + # those contributions — we pay the wasted GEMM compute because batched_mm has no offset to skip. + expert_ids.clamp_(0, self.num_experts - 1) + + # Select gate_up or just up projection weights and biases + if self.has_gate: + selected_weights = self.gate_up_proj[expert_ids] + selected_biases = self.gate_up_proj_bias[expert_ids] if self.has_bias else None + else: + selected_weights = self.up_proj[expert_ids] + selected_biases = self.up_proj_bias[expert_ids] if self.has_bias else None + + # --- Up projection per expert (batched) --- + proj_out = _batched_linear( + selected_hidden_states, selected_weights, bias=selected_biases, is_transposed=self.is_transposed + ) # (S, 2 * intermediate_dim) or (S, intermediate_dim) depending on whether we have gating + + # Apply gating or activation + if self.has_gate: + # for gated experts we apply the custom/default gating mechanism + proj_out = self._apply_gate(proj_out) # (S, intermediate_dim) + else: + # for non-gated experts we just apply the activation function + proj_out = self.act_fn(proj_out) # (S, intermediate_dim) + + # Select down projection weights and biases for selected samples + selected_weights = self.down_proj[expert_ids] + selected_biases = self.down_proj_bias[expert_ids] if self.has_bias else None + + # --- Down projection per expert (batched) --- + proj_out = _batched_linear( + proj_out, selected_weights, bias=selected_biases, is_transposed=self.is_transposed + ) # (S, hidden_dim) + + # Apply routing weights + weighted_out = proj_out * sample_weights.unsqueeze(-1) # (S, hidden_dim) + + # Accumulate results using deterministic reshape+sum instead of index_add_ + # index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd + # index_add_ accumulates in-place using the dtype of the output tensor (fp16/bf16) + # reshape+sum accumulates in fp32 which is more stable for low precision training/inference. + final_hidden_states = weighted_out.view(num_tokens, num_top_k, hidden_dim).sum(dim=1) + + return final_hidden_states.to(hidden_states.dtype) + + +# torch.compiler.disable does not work with fullgraph=True, so we implement a custom operator to opaque this function. +# This is not "free compilation compatibility" because now inductor won't be able to optimize matmuls inside the loop, +# but since the matmuls here have dynamic shapes, inductor wouldn't have been able to optimize them anyway. +def _grouped_mm_fallback(input: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor) -> torch.Tensor: + """ + Fallback grouped matrix multiplication used when `torch.nn.functional.grouped_mm` and `torch._grouped_mm` + are unavailable or incompatible with `torch.compile` (e.g. non-bfloat16 weights). + + Args: + input (`torch.Tensor`): Input of shape (S, input_dim), sorted by expert id. + weight (`torch.Tensor`): Expert weights of shape (num_experts, input_dim, output_dim). + offs (`torch.Tensor`): Cumulative token counts per expert of shape (num_experts,). + Returns: + `torch.Tensor`: Output of shape (S, output_dim). + """ + output = torch.zeros(input.size(0), weight.size(2), device=input.device, dtype=input.dtype) # (S, output_dim) + + start = 0 + # single cpu<->gpu sync point here, + # avoids multiple syncs inside the loop + for i, end in enumerate(offs.tolist()): + if start == end: + continue + torch.mm(input[start:end], weight[i], out=output[start:end]) + start = end + + return output + + +def _grouped_mm_fallback_fake(input: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor) -> torch.Tensor: + """Shape/dtype inference stub for `_grouped_mm_fallback` required by `torch.compile`.""" + assert input.dim() == 2, f"input must be 2D (S, input_dim), got shape {tuple(input.shape)}" + assert weight.dim() == 3, ( + f"weight must be 3D (num_experts, input_dim, output_dim), got shape {tuple(weight.shape)}" + ) + assert offs.dim() == 1, f"offs must be 1D (num_experts,), got shape {tuple(offs.shape)}" + assert offs.size(0) == weight.size(0), f"offs length {offs.size(0)} must match number of experts {weight.size(0)}" + assert input.size(1) == weight.size(1), ( + f"input_dim mismatch: input has {input.size(1)}, weight has {weight.size(1)}" + ) + assert offs.dtype in (torch.int32, torch.int64), f"offs must be an integer tensor, got {offs.dtype}" + return torch.empty(input.size(0), weight.size(2), device=input.device, dtype=input.dtype) + + +def _grouped_mm_fallback_setup_context(ctx, inputs, output): + """Saves input and weight for backward; offs is stored directly as it is a non-differentiable integer tensor.""" + ctx.save_for_backward(inputs[0], inputs[1]) + ctx.offs = inputs[2] + + +def _grouped_mm_fallback_backward(ctx, grad_output): + """Backward pass for `_grouped_mm_fallback`. Computes grad_input and grad_weight per expert group; offs has no gradient.""" + input, weight = ctx.saved_tensors + grad_input = torch.zeros_like(input) + grad_weight = torch.zeros_like(weight) + + start = 0 + # single cpu<->gpu sync point here, + # avoids multiple syncs inside the loop + for i, end in enumerate(ctx.offs.tolist()): + if start == end: + continue + torch.mm(grad_output[start:end], weight[i].T, out=grad_input[start:end]) + torch.mm(input[start:end].T, grad_output[start:end], out=grad_weight[i]) + start = end + + return grad_input, grad_weight, None + + +if is_torch_available(): + torch.library.custom_op( + "transformers::grouped_mm_fallback", + _grouped_mm_fallback, + mutates_args=(), + schema="(Tensor input, Tensor weight, Tensor offs) -> Tensor", + ) + torch.library.register_fake("transformers::grouped_mm_fallback", _grouped_mm_fallback_fake) + torch.library.register_autograd( + "transformers::grouped_mm_fallback", + _grouped_mm_fallback_backward, + setup_context=_grouped_mm_fallback_setup_context, + ) + + +def _can_use_grouped_mm(input: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor) -> bool: + """ + Check if torch.nn.functional.grouped_mm or torch._grouped_mm can be used based on availability and compatibility with torch.compile. + + Args: + input (`torch.Tensor`): + Input tensor of shape (S, input_dim). + weight (`torch.Tensor`): + Weight tensor of shape (num_experts, input_dim, output_dim). + offs (`torch.Tensor`): + Offsets tensor indicating the boundaries of each group in the input tensor. + Returns: + `bool`: True if grouped_mm can be used, False otherwise. + """ + if (is_torchdynamo_compiling() and weight.dtype != torch.bfloat16) or ( + weight.device.type == "cpu" + # accept_dev=True is necessary for "+cpu"/"+xpu" etc. + and is_torch_less_or_equal("2.10.0", accept_dev=True) + and (weight.data_ptr() % 16 != 0 or input.data_ptr() % 16 != 0) + ): + # Under the following conditions we cannot use torch.grouped_mm and have to fall back: + # 1. torch.grouped_mm is not supported in torch.compile / inductor with dtypes other than bf16 + # 2. Before PyTorch 2.11, torch.grouped_mm on CPU required 16 bytes alignment which is not + # guaranteed for tensors loaded using memmap (e.g. using safetensors lazy tensor loading) + # and not really necessary because the cpu path uses a fallback for-loop implementation. + # issue: https://github.com/pytorch/pytorch/issues/172440 + return False + + # On CUDA, `grouped_mm` availability also depends on GPU compute capability: + # `torch.nn.functional.grouped_mm` in torch>=2.10 and `torch._grouped_mm` in torch>=2.9 support SM80+ + # but older `torch._grouped_mm` requires SM90+. + if weight.device.type == "cuda": + if hasattr(torch.nn.functional, "grouped_mm"): + return torch.cuda.get_device_capability(weight.device) >= (8, 0) + if hasattr(torch, "_grouped_mm"): + if is_torch_greater_or_equal("2.9", accept_dev=True): + return torch.cuda.get_device_capability(weight.device) >= (8, 0) + else: + return torch.cuda.get_device_capability(weight.device) >= (9, 0) + + return False + + return hasattr(torch.nn.functional, "grouped_mm") or hasattr(torch, "_grouped_mm") + + +def _grouped_mm( + input: torch.Tensor, + weight: torch.Tensor, + offs: torch.Tensor, +) -> torch.Tensor: + """Grouped matrix multiplication dispatcher that uses torch.nn.functional.grouped_mm if available, else falls back to torch._grouped_mm. + + Args: + input (`torch.Tensor`): + Input tensor of shape (S, input_dim). + weight (`torch.Tensor`): + Weight tensor of shape (num_experts, input_dim, output_dim). + offs (`torch.Tensor`): + Offsets tensor indicating the boundaries of each group in the input tensor. + Returns: + `torch.Tensor`: Output tensor of shape (S, output_dim). + """ + + if _can_use_grouped_mm(input, weight, offs): + # torch.nn.functional.grouped_mm and torch._grouped_mm are not autocast-enabled, + # when autocast is enabled we can end up with intermediate tensors in fp32 (e.g. LayerNorm output) and weight tensors in bf16 + # In that case we need to cast the input to the weight dtype to avoid dtype mismatch errors. + # See: https://github.com/pytorch/pytorch/issues/174763 + if hasattr(torch.nn.functional, "grouped_mm"): + return torch.nn.functional.grouped_mm(input.to(weight.dtype), weight, offs=offs) + elif hasattr(torch, "_grouped_mm"): + return torch._grouped_mm(input.to(weight.dtype), weight, offs=offs) + + return torch.ops.transformers.grouped_mm_fallback(input, weight, offs=offs) + + +def _grouped_linear( + input: torch.Tensor, + weight: torch.Tensor, + offs: torch.Tensor, + bias: torch.Tensor | None = None, + is_transposed: bool = False, +) -> torch.Tensor: + """Grouped linear layer supporting optional bias and transposed weights. + + Args: + input (`torch.Tensor`): + Input tensor of shape (S, input_dim). + weight (`torch.Tensor`): + Weight tensor of shape (num_experts, input_dim, output_dim) if `is_transposed`, + else of shape (num_experts, output_dim, input_dim). + offs (`torch.Tensor`): + Offsets tensor indicating the boundaries of each group in the input tensor. + bias (`torch.Tensor`, *optional*): + Bias tensor of shape (num_experts, output_dim). Default is `None`. + is_transposed (`bool`, *optional*, defaults to `False`): + Whether the weight tensor is transposed. + Returns: + `torch.Tensor`: Output tensor of shape (S, output_dim). + """ + if is_transposed: + # (S, input_dim) @ grouped (num_experts, input_dim, output_dim) -> (S, output_dim) + out = _grouped_mm(input, weight, offs=offs) + else: + # (S, input_dim) @ grouped (num_experts, output_dim, input_dim).T -> (S, output_dim) + out = _grouped_mm(input, weight.transpose(-2, -1), offs=offs) + + if bias is not None: + # We should be able to pass bias to the grouped_mm call, but it's not yet supported. + out.add_(bias) + + return out + + +def grouped_mm_experts_forward( + self: torch.nn.Module, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + device = hidden_states.device + num_top_k = top_k_index.size(-1) + num_tokens = hidden_states.size(0) + hidden_dim = hidden_states.size(-1) + + # S is the number of selected tokens-experts pairs (S = num_tokens * num_top_k) + sample_weights = top_k_weights.reshape(-1) # (S,) + expert_ids = top_k_index.reshape(-1) # (S,) + + # Sort by expert for grouped processing + expert_ids_g, perm = torch.sort(expert_ids) + selected_hidden_states_g = hidden_states[perm // num_top_k] + sample_weights_g = sample_weights[perm] + + # Compute offsets for grouped_mm + # using histc instead of bincount to avoid cuda graph issues + # With deterministic algorithms, CPU only supports float input, CUDA only supports int input. + # torch.histc() does not support integer dtypes on CPU and MPS. + histc_input = expert_ids_g.float() if device.type in ("cpu", "mps") else expert_ids_g.int() + tokens_per_expert = torch.histc(histc_input, bins=self.num_experts, min=0, max=self.num_experts - 1) + offsets = torch.cumsum(tokens_per_expert, dim=0, dtype=torch.int32) + + # EP sentinel handling: leave `expert_ids` unclamped so the sort pushes sentinels to the tail, + # `histc(max=num_experts-1)` drops them from `tokens_per_expert`, and grouped_mm skips rows + # beyond `offsets[-1]` — sentinels cost no real GEMM compute. The kernel leaves sentinel-tail + # rows of its output uninit (both fwd output and bwd `d_input`), but ONE pre-mask + ONE + # post-mask covers the whole forward — no per-grouped_mm masking is needed, because + # intermediate sentinel-row NaN is only ever consumed by the next grouped_mm, which itself + # only reads rows `< offsets[-1]`: + # - fwd post-mask on `weighted_out`: kills `proj_out[sentinel] * 0 = NaN * 0 = NaN` + # before the per-token reduction sums it. + # - bwd pre-mask on `selected_hidden_states_g`: its `masked_fill_` backward zeros sentinel + # rows of `d_selected_hidden_states_g` after the up grouped_mm bwd writes them as + # uninit, and before the gather's scatter-add pushes them into `d_hidden_states`. + # In-place clamp on `expert_ids_g` keeps the per-row bias gather in-bounds (bias added at + # sentinel positions falls in rows the kernel skips, so harmless). Safe to mutate now — + # nothing downstream needs the sentinel info from `expert_ids_g` itself. + sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1) + expert_ids_g.clamp_(max=self.num_experts - 1) + + # Select expert weights and biases + # NOTE: We keep all experts here and rely on offsets to target the active ones. + # I have already implemented a version that only passes the active experts, but + # to do so I had to use torch.unique which breaks the graph capture (data-dependent). + # Also there were no speedup gains from it in my experiments, even in eager mode. + # NOTE: The grouped_mm kernel only targets the active experts / tokens via the offsets + if self.has_gate: + selected_weights = self.gate_up_proj + selected_biases = self.gate_up_proj_bias[expert_ids_g] if self.has_bias else None + else: + selected_weights = self.up_proj + selected_biases = self.up_proj_bias[expert_ids_g] if self.has_bias else None + + # Pre-mask (bwd path). + selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0) + + # --- Up projection per expert (grouped) --- + proj_out = _grouped_linear( + selected_hidden_states_g, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed + ) # (S, 2 * intermediate_dim) or (S, intermediate_dim) depending on whether we have gating + + # Apply gating or activation + if self.has_gate: + # for gated experts we apply the custom/default gating mechanism + proj_out = self._apply_gate(proj_out) # (S, intermediate_dim) + else: + # for non-gated experts we just apply the activation function + proj_out = self.act_fn(proj_out) # (S, intermediate_dim) + + # Select down projection weights and biases + selected_weights = self.down_proj + selected_biases = self.down_proj_bias[expert_ids_g] if self.has_bias else None + + # --- Down projection per expert (grouped) --- + proj_out = _grouped_linear( + proj_out, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed + ) # (S, hidden_dim) + + # Apply routing weights + weighted_out = proj_out * sample_weights_g.unsqueeze(-1) # (S, hidden_dim) + + # Post-mask (fwd path). + weighted_out.masked_fill_(sentinel_mask, 0.0) + + # Restore original order + inv_perm = torch.empty_like(perm) + inv_perm[perm] = torch.arange(perm.size(0), device=device) + weighted_out = weighted_out[inv_perm] # (S, hidden_dim) + + # Accumulate results using deterministic reshape+sum instead of index_add_ + # index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd + # index_add_ accumulates in-place using the dtype of the output tensor (fp16/bf16) + # reshape+sum accumulates in fp32 which is more stable for low precision training/inference. + final_hidden_states = weighted_out.view(num_tokens, num_top_k, hidden_dim).sum(dim=1) + + return final_hidden_states.to(hidden_states.dtype) + + +class ExpertsInterface(GeneralInterface): + """Interface for registering custom experts forward functions.""" + + _global_mapping = { + "batched_mm": batched_mm_experts_forward, + "grouped_mm": grouped_mm_experts_forward, + "sonicmoe": sonicmoe_experts_forward, + } + + def get_interface(self, experts_implementation: str, default: Callable) -> Callable: + """Return the requested `experts_implementation`. Also strictly check its validity, and raise if invalid.""" + if experts_implementation is None: + logger.warning_once( + "You tried to access the `ExpertsInterface` with a `config._experts_implementation` set to `None`. This " + "is expected if you use an Expert Module as a standalone Module. If this is not the case, something went " + "wrong with the dispatch of `config._experts_implementation`" + ) + elif experts_implementation != "eager" and experts_implementation not in self: + raise KeyError( + f"`{experts_implementation}` is not a valid experts implementation registered in the `ExpertsInterface`" + ) + return super().get(experts_implementation, default) + + +ALL_EXPERTS_FUNCTIONS = ExpertsInterface() + + +def _default_apply_gate(self, gate_up_out: torch.Tensor) -> torch.Tensor: + """ + Default gating mechanism: splits the gate_up_out into gate and up parts, + applies the activation function to the gate part, and multiplies it with the up part. + Args: + gate_up_out (`torch.Tensor`): + The output tensor from the gate and up projection of shape (S, 2 * intermediate_dim). + Returns: + `torch.Tensor`: The gated output tensor of shape (S, intermediate_dim). + """ + gate, up = gate_up_out.chunk(2, dim=-1) # (S, intermediate_dim) + return self.act_fn(gate) * up # (S, intermediate_dim) + + +def use_experts_implementation( + experts_class: type[torch.nn.Module] | None = None, + *, + experts_interface: ExpertsInterface = ALL_EXPERTS_FUNCTIONS, + is_concatenated: bool = True, + is_transposed: bool = False, + has_bias: bool = False, + has_gate: bool = True, +) -> type[torch.nn.Module]: + """Decorator to modify experts class to support different experts implementations. + + Args: + experts_class (`type[torch.nn.Module]`, *optional*): + The experts class to modify. If not provided, returns a decorator that can be applied to the class. + experts_interface (`ExpertsInterface`, *optional*, defaults to `ALL_EXPERTS_FUNCTIONS`): + The experts interface to use for dispatching the forward method. + is_concatenated (`bool`, *optional*, defaults to `True`): + Whether the expert weights are stored in concatenated layout [gate;up] + or interleaved layout [gate0, up0, gate1, up1, ...]. + is_transposed (`bool`, *optional*, defaults to `False`): + Whether the expert weights are stored in transposed format. + has_bias (`bool`, *optional*, defaults to `False`): + Whether the expert layers include bias terms or not. + has_gate (`bool`, *optional*, defaults to `True`): + Whether the experts use a gating mechanism or not. + Whether it has gate_up_proj weights or just up_proj weights. + + Returns: + `type[torch.nn.Module]`: The modified experts class. + """ + + def wrapper(experts_class: type[torch.nn.Module]) -> type[torch.nn.Module]: + original_init = experts_class.__init__ + original_forward = experts_class.forward + + @wraps(original_init) + def __init__(self, config, *args, **kwargs): + original_init(self, config, *args, **kwargs) + self.config = config + self.has_gate = has_gate + self.has_bias = has_bias + self.is_transposed = is_transposed + self.is_concatenated = is_concatenated + + @wraps(original_forward) + def forward(self, *args, **kwargs): + experts_forward = experts_interface.get_interface(self.config._experts_implementation, original_forward) + return experts_forward(self, *args, **kwargs) + + if not hasattr(experts_class, "_apply_gate"): + experts_class._apply_gate = _default_apply_gate + + experts_class.__init__ = __init__ + experts_class.forward = forward + return experts_class + + if experts_class is not None: + return wrapper(experts_class) + + return wrapper diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/npu_flash_attention.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/npu_flash_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..2af3ba84e012c88224e339886309379ebd2bf8cb --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/npu_flash_attention.py @@ -0,0 +1,143 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import os + +import torch + +from ..utils.import_utils import is_torch_npu_available + + +if is_torch_npu_available(): + from torch_npu import npu_fusion_attention + + +# FlashAttention2 is supported on Ascend NPU with down-right aligned causal mask by default. +# Set environment variable `NPU_FA2_SPARSE_MODE` to 2 when using top-left aligned causal mask. +TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE = 2 +DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE = 3 + +SPARSE_MODE = int(os.getenv("NPU_FA2_SPARSE_MODE", default=DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE)) +if SPARSE_MODE not in [TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE, DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE]: + raise ValueError( + "Environment variable `NPU_FA2_SPARSE_MODE` can only be set as 2 (top-left aligned causal mask) " + "or 3 (down-right aligned causal mask)." + ) + +ATTN_MASK_NPU_CACHE = {} + + +def get_attn_mask_npu(device): + """Get or create attention mask for the specified device.""" + if device not in ATTN_MASK_NPU_CACHE: + ATTN_MASK_NPU_CACHE[device] = torch.triu(torch.ones([2048, 2048], device=device), diagonal=1).bool() + return ATTN_MASK_NPU_CACHE[device] + + +def is_npu_fa2_top_left_aligned_causal_mask(): + return SPARSE_MODE == TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE if is_torch_npu_available() else False + + +def npu_flash_attn_func( + q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=False, + **kwargs, +): + keep_prob = 1.0 - dropout_p + + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(q.shape[-1]) + + if not causal: + head_num = q.shape[2] + output = npu_fusion_attention(q, k, v, head_num, "BSND", keep_prob=keep_prob, scale=softmax_scale)[0] + else: + attn_mask_npu = get_attn_mask_npu(q.device) + head_num = q.shape[2] + output = npu_fusion_attention( + q, + k, + v, + head_num, + "BSND", + keep_prob=keep_prob, + scale=softmax_scale, + atten_mask=attn_mask_npu, + sparse_mode=SPARSE_MODE, + )[0] + + return output + + +def npu_flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q=None, # defined for aligning params order with corresponding function in `flash-attn` + max_seqlen_k=None, # defined for aligning params order with corresponding function in `flash-attn` + dropout_p=0.0, + softmax_scale=None, + causal=False, + **kwargs, +): + keep_prob = 1.0 - dropout_p + + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(q.shape[-1]) + + if not causal: + head_num = q.shape[1] + output = npu_fusion_attention( + q, + k, + v, + head_num, + pse=None, + atten_mask=None, + scale=softmax_scale, + keep_prob=keep_prob, + input_layout="TND", + actual_seq_qlen=tuple(cu_seqlens_q[1:].cpu().numpy().tolist()), + actual_seq_kvlen=tuple(cu_seqlens_k[1:].cpu().numpy().tolist()), + )[0] + else: + attn_mask_npu = get_attn_mask_npu(q.device) + head_num = q.shape[1] + output = npu_fusion_attention( + q, + k, + v, + head_num, + pse=None, + padding_mask=None, + atten_mask=attn_mask_npu, + scale=softmax_scale, + keep_prob=keep_prob, + input_layout="TND", + actual_seq_qlen=tuple(cu_seqlens_q[1:].cpu().numpy().tolist()), + actual_seq_kvlen=tuple(cu_seqlens_k[1:].cpu().numpy().tolist()), + sparse_mode=SPARSE_MODE, + )[0] + + return output + + +# This function is not implemented but should never be called because block table is not used on NPU +def npu_flash_attn_with_kvcache(): + raise NotImplementedError("npu_flash_attn_with_kvcache is not implemented") diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/spqr.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/spqr.py new file mode 100644 index 0000000000000000000000000000000000000000..d7b2ae7d3ab1b328755c469063a9284d78cb7e62 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/spqr.py @@ -0,0 +1,74 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"SpQR (Sparse-Quantized Representation) integration file" + +from ..quantizers.quantizers_utils import should_convert_module +from ..utils import is_spqr_available, is_torch_available, logging + + +if is_torch_available(): + import torch + import torch.nn as nn + +logger = logging.get_logger(__name__) + + +def replace_with_spqr_linear(model, modules_to_not_convert: list[str] | None = None, quantization_config=None): + """ + Public method that replaces the Linear layers of the given model with SPQR quantized layers. + + Args: + model (`torch.nn.Module`): + The model to convert, can be any `torch.nn.Module` instance. + modules_to_not_convert (`list[str]`, *optional*, defaults to `None`): + A list of nn.Linear weights to not convert. If a parameter path is in the list (e.g. `lm_head.weight`), the corresponding module will not be + converted. + quantization_config (`SpQRConfig`): + The quantization config object that contains the quantization parameters. + """ + if is_spqr_available(): + from spqr_quant import QuantizedLinear + + has_been_replaced = False + # we need this to correctly materialize the weights during quantization + for module_name, module in model.named_modules(): + if not should_convert_module(module_name, modules_to_not_convert): + continue + with torch.device("meta"): + if isinstance(module, nn.Linear): + shapes = quantization_config.shapes + + new_module = QuantizedLinear.create_placehodler( + rows=module.out_features, + cols=module.in_features, + bits=quantization_config.bits, + beta1=quantization_config.beta1, + beta2=quantization_config.beta2, + dense_weights_shape=shapes[f"{module_name}.dense_weights.shape"], + row_offsets_shape=shapes[f"{module_name}.row_offsets.shape"], + col_vals_shape=shapes[f"{module_name}.col_vals.shape"], + in_perm_shape=shapes[f"{module_name}.in_perm.shape"], + ) + # Force requires grad to False to avoid unexpected errors + model._modules[module_name].requires_grad_(False) + model.set_submodule(module_name, new_module) + has_been_replaced = True + if not has_been_replaced: + logger.warning( + "You are loading your model using eetq but no linear modules were found in your model." + " Please double check your model architecture, or submit an issue on github if you think this is" + " a bug." + ) + + return model diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/torchao.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/torchao.py new file mode 100644 index 0000000000000000000000000000000000000000..421a004dd6e924a199d7d8cb96eb3d4fc6d4294a --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/integrations/torchao.py @@ -0,0 +1,234 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +import types + +import torch + +from transformers.utils import logging +from transformers.utils.import_utils import is_torch_available, is_torchao_available + + +if is_torch_available(): + from ..core_model_loading import ConversionOps +from ..quantizers.quantizers_utils import get_module_from_name + + +if is_torchao_available(): + from torchao.prototype.safetensors.safetensors_support import ( + unflatten_tensor_state_dict, + ) + from torchao.prototype.safetensors.safetensors_utils import is_metadata_torchao + +logger = logging.get_logger(__name__) + + +def _quantization_type(weight): + from torchao.dtypes import AffineQuantizedTensor + from torchao.quantization.linear_activation_quantized_tensor import LinearActivationQuantizedTensor + + if isinstance(weight, AffineQuantizedTensor): + return f"{weight.__class__.__name__}({weight._quantization_type()})" + + if isinstance(weight, LinearActivationQuantizedTensor): + return f"{weight.__class__.__name__}(activation={weight.input_quant_func}, weight={_quantization_type(weight.original_weight_tensor)})" + + +def _linear_extra_repr(self): + weight = _quantization_type(self.weight) + if weight is None: + return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight=None" + else: + return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight={weight}" + + +class TorchAoQuantize(ConversionOps): + def __init__(self, hf_quantizer): + self.hf_quantizer = hf_quantizer + + def _quantize(self, module, config, *args, **kwargs): + """Run quantize_, moving to CUDA first if CPU offloading is active. + + Some torchao quantization ops (e.g. int4 packing) only have CUDA kernels. + When a layer is destined for CPU (e.g. CPU offloading), we temporarily move + it to CUDA for quantization, then move the result back to CPU. + """ + from torchao.quantization import quantize_ + + target_device = next(module.parameters()).device + if self.hf_quantizer.offload_to_cpu and target_device.type == "cpu": + module.to("cuda") + quantize_(module, config, *args, **kwargs) + module.to("cpu") + else: + quantize_(module, config, *args, **kwargs) + + def convert( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module | None = None, + full_layer_name: str | None = None, + missing_keys=None, + **kwargs, + ) -> dict[str, torch.Tensor]: + _, value = tuple(input_dict.items())[0] + value = value[0] if isinstance(value, list) else value + + module, tensor_name = get_module_from_name(model, full_layer_name) + + module._parameters[tensor_name] = torch.nn.Parameter(value, requires_grad=value.requires_grad) + # if we are quantizing tied parameters, to avoid tying the quantized weights + # the correct order to do it is + # 1. load the weight to model + # 2. run tie_weights to populate the weights + # 3. quantize + input_embed = model.get_input_embeddings() + is_embedding_param = id(module) == id(input_embed) + untie_embedding_weights = self.hf_quantizer.quantization_config.untie_embedding_weights + + if untie_embedding_weights and is_embedding_param: + setattr(model.config.get_text_config(decoder=True), "tie_word_embeddings", False) + + from torchao.quantization import FqnToConfig + + config = self.hf_quantizer.quantization_config.get_apply_tensor_subclass() + if isinstance(config, FqnToConfig): + module_fqn, top_level_param_name = full_layer_name.rsplit(".", 1) + c = None + if full_layer_name in config.fqn_to_config: + assert not module_fqn.startswith("re:"), ( + "param fqn should not start with`re:`, which is used for specifying regex" + ) + c = config.module_fqn_to_config[full_layer_name] + elif module_fqn in config.fqn_to_config: + assert not module_fqn.startswith("re:"), ( + "module fqn should not start with`re:`, which is used for specifying regex" + ) + c = config.module_fqn_to_config[module_fqn] + # regex match module and param + else: + for maybe_module_fqn_pattern in config.fqn_to_config: + # if key doesn't start with re, it is an exact fqn key, so we don't regex match + if not maybe_module_fqn_pattern.startswith("re:"): + continue + # see if param matches first + elif re.fullmatch(maybe_module_fqn_pattern[3:], full_layer_name): + c = config.module_fqn_to_config[maybe_module_fqn_pattern] + break + elif re.fullmatch(maybe_module_fqn_pattern[3:], module_fqn): + # we'll apply the config for first fully matched pattern + c = config.module_fqn_to_config[maybe_module_fqn_pattern] + break + else: + c = config.module_fqn_to_config.get("_default", None) + + if c is not None: + if top_level_param_name == "weight": + if is_embedding_param and untie_embedding_weights: + lm_head = module.weight.clone() + # we can apply the module config directly + self._quantize(module, c, (lambda x, fqn: True)) + missing_keys.discard(full_layer_name) + module._is_hf_initialized = True + # torchao quantizes weights into a module but some models access the weight directly + # (e.g. module.o_proj.weight). The _is_hf_initialized flag is set at the module + # level only, so we also set it on each parameter to prevent _init_weights from + # calling normal_() on already-quantized Float8Tensors. + for param in module.parameters(recurse=False): + param._is_hf_initialized = True + return {"lm_head.weight": lm_head} if is_embedding_param and untie_embedding_weights else {} + else: + # need to apply to custom param name + custom_param_fqn_config = FqnToConfig({top_level_param_name: c}) + self._quantize(module, custom_param_fqn_config, filter_fn=None) + missing_keys.discard(full_layer_name) + module._is_hf_initialized = True + for param in module.parameters(recurse=False): + param._is_hf_initialized = True + return {} + return {full_layer_name: value} + + if is_embedding_param and untie_embedding_weights: + lm_head = module.weight.clone() + self._quantize(module, self.hf_quantizer.quantization_config.get_apply_tensor_subclass()) + missing_keys.discard(full_layer_name) + module._is_hf_initialized = True + for param in module.parameters(recurse=False): + param._is_hf_initialized = True + return {"lm_head.weight": lm_head} if is_embedding_param and untie_embedding_weights else {} + + +class TorchAoDeserialize(ConversionOps): + def __init__(self, hf_quantizer): + self.hf_quantizer = hf_quantizer + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str] | None = None, + model: torch.nn.Module | None = None, + full_layer_name: str | None = None, + missing_keys=None, + **kwargs, + ) -> dict[str, torch.Tensor]: + """ + Consolidates tensor subclass components before reconstructing the object + + For example: + input_dict: { + "_weight_qdata": torch.Tensor, + "_weight_scale": torch.Tensor, + } + full_layer_name: "model.layers.0.self_attn.k_proj.weight" + + Given this, we reconstruct a Float8Tensor instance using the qdata and scale + and return it as a dictionary with the full_layer_name as the key and the recovered + Float8Tensor instance as the value. + """ + is_unsafe_serialization = list(input_dict.keys())[0] not in source_patterns + + param_data = {} + layer_name = ".".join(full_layer_name.split(".")[:-1]) + if is_unsafe_serialization: + if isinstance(input_dict["weight"], list): + weight = input_dict["weight"][0] + else: + weight = input_dict["weight"] + else: + for suffix in input_dict.keys(): + if len(input_dict[suffix]) != 1: + raise ValueError( + f"Expected a single tensor for {suffix} but got {len(input_dict[suffix])} tensors instead" + ) + param_data[f"{layer_name}.{suffix}"] = input_dict[suffix][0] + + # If it's unsafe-serialized (i.e. not safetensors), no need for anything + if is_unsafe_serialization: + return {full_layer_name: weight} + elif not is_metadata_torchao(self.hf_quantizer.metadata): + raise ValueError("Invalid torchao safetensors metadata") + + unflattened_state_dict, leftover_state_dict = unflatten_tensor_state_dict( + param_data, self.hf_quantizer.metadata + ) + assert not leftover_state_dict # there should be no unprocessed tensors + new_param = unflattened_state_dict[full_layer_name] + + module, _ = get_module_from_name(model, full_layer_name) + # Add repr to the module + if isinstance(module, torch.nn.Linear): + module.extra_repr = types.MethodType(_linear_extra_repr, module) + + return {full_layer_name: new_param} diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/logs/infer_owt_llmclean_qwen36_35b_10k_C1to1024sqrt_step013000_gpu_temp1_decode128_quick_n8.log b/LTA_openwebtext_dualt/mini_owt_logdirichlet/logs/infer_owt_llmclean_qwen36_35b_10k_C1to1024sqrt_step013000_gpu_temp1_decode128_quick_n8.log new file mode 100644 index 0000000000000000000000000000000000000000..069c6f36bf3cb59e6a4fa5e5dc4aa11b15b9a9c2 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/logs/infer_owt_llmclean_qwen36_35b_10k_C1to1024sqrt_step013000_gpu_temp1_decode128_quick_n8.log @@ -0,0 +1,34 @@ +bos=1: eos=1: +decode steps=128 temp=1.0 mode=dirichlet c_min=1.0 c_max=1024.0 c_schedule=sqrt train_c_min=1.0 train_c_max=1024.0 +===== sample 0 ===== +head_tokens: ['ing', '▁of', '▁the', '▁conviction', '▁of', '▁the', '▁capital', '▁of', '▁', 'a', 'FS', '▁of', '▁the', '▁returns', '▁of', '▁the'] +tail_tokens: ['▁and', '▁occurred', '▁to', '▁the', '▁the', 's', '▁of', '▁the', '▁understanding', '▁of', '▁the', '▁triumph', '▁of', '▁the', '▁central', ''] +ing of the conviction of the capital of aFS of the returns of the god of the statuees, in the miracle of the capital. From up to pass the understanding of the Spiritss of the shields of the substance of the doors of thewards. Following the background of the initial account of contribution to the expression of the consciousness of the principal portion of their relations. principal collapses the remains in the refusal to ounce the ins of the frequentguide of Treasure. The understanding of the original account of the mind presented on the formation of the part of account the account of the understanding of theinals of the truth of the dis approval of the revolution, and the account of the Conquis of the the account of the cane of the authority. This is the lits of the property and the the narrative of theuring advance of the death the prominents of theigiken in the ambassador of the landlord of close in the age of the141. The noble Eli seeks to the coaces of the70 from the deepens of the authority of the law as the only remains of the gods and question of the actions in the ens of authority of the connection of the Royalativesment of the noises in thees possession of the entireclave of the crown of the Constitution to the relation of the head of the evidence of property in the Brazilian Court of the word created in the current corner of the value of the gate, thee of the power of the dealings of the believers of the the value of the Holyclaver of theors of detail, which that is the sound of the lored out of the Book of the, and the sources of the predicts of the essential of the expected to account the formation of the truth and by the truth of the property. According to hear the account of the commissioner contains as the referredly from the leadership of the battle of the ongoing authority of the expansion, to apply thecane of the battle to the understanding of the concrate, the understanding of the question of the ennics the ease of the truth and the constants of the power, and index the legacy of the law of the concentrations. The organization of the destruction of the fake faithful and the victim of the control of the parishs of theward itself of the site of the Conramas of the 13mination of the intervention, of the sens from the account of the Book of the unit. By the truth of the pastor of ten of the victim of the dus of the correct, ounce the equals of the Church of the theoretical theories of the understanding of spirituals, to thee of the essential eyes. According in the account of the Revolution of the transmission of the Christ of the relations of the gate of the first and account of the end of the occasion and the citizen and scopes on the top of the exhibit of societys and the remains of the truth of that of theclave of the 300-iration of the grand gate of the thro into a constant order. The understanding of the lack of the property, to the efficiency of the crown of thees of the doors to theimposed of the legacy of the property. The superior to the center of the contents of the the hands of the hands of the imposed of the legacy. It is the official explanation70 of the formation of the category and investigations upon theeon of the pen Church of the sum of the State of the epicill Christ. The Crowns a detailly as a particular truth of the constants of explanation. Just in the position of the reporting of the wards in the Vics the representatives of the grandes of the state, without aative of the category of the parishes of the card of the head of the power of the order of the Popes, in the display. Following the understanding of the writer and the hands of the priorly of thecan of the translation and truth of the page. The pre contribution from the extent of consciousness of the translations of the Caesars of the understanding and the count of the date of truth is, as the power of the viative of the truth of the moment of the mind of the wrong truth is a matter of practice in the contribution of the death of the discovery, in the relation of the noble, and the reflection of the dispower. The understanding Ministry of the will of the exchange of the saative of the Churchs of the entertainment. The truth is the understanding of the power of the forces of the situation in. It is the means of the account, and upon the abolishes in the Book of the world of a teacher to present in the remains of the significant understandings of the power. To the control of the Constitution of s the subject of the treatment of the law account and occurred to the thes of the understanding of the triumph of the central +===== sample 1 ===== +head_tokens: ['▁Roger', 's', '.', '▁I', '▁was', '▁convinced', '▁that', '▁', 'he', "'", 's', '▁been', '▁as', '▁the', '▁writing', '▁team'] +tail_tokens: ["'", 's', '▁eyes', ',', '▁', 'a', '▁show', '▁or', '▁', 'a', 'head', '▁to', '▁him', '.', '▁I', ''] +Rogers. I was convinced that he's been as the writing team, and he's getting to do him as understand that he's the officer a challenge to his obligations. "What a year, and he's a nice VI, and that's a a a sout conclusion. In some way, the practitioner: staff has call a heart of the weight. "I, it's the ssies, he's thought he has a Defense generation, no person.. I'm preparing a second person, um a hos. But it is a bit of a, the man drawn of the way, he's a bit. So to that, but he's told, he says his wife, a guy, a fit of the film, to raise him, he's it's a great. I'm a bunch. That's a trend on, and he's not to give a utters. But, that? Yeah, is he's, so. 'Some of all his stance, he's, and he's a fan of his eye funny. He's like he's the likes, he's a name, it's a9,' he grandmother's a quick, in the creative. Sometimes he's in a utter guy, he has him, of his life, he, he's like a finds. I a rathere, maybe, he's a guy I tell him to be watching. Sometimes that's him, I don't tell him a teenager call, he's him, that's absolutely tene's quiger.'" In a way that's me, he's a teenager, he's a Obster, he 41, he's very silent. His goence, he's done it, occasion, a app that's a on he's been in a quarter of politics, in a growing state, he still's a lack of doubt, he's a call to a-on, or a importance level of sort of dark and the stene of 40 and, he's getting in a way, he's there. He says a trend he says he's not in. he's got, a personality story from the zform of Fall. That's things he's a case, he's nice, he's a man a really narrow. This is a voicement. A sort of decades in a sea-auist, and he still's a ard in a fair call: and he's if he's he's, yeah. I's. I got a teenagers, 12, Imes. I had decided that shot him, he's a call with a a ant, he talks to that, he's, always see it, he's thoughtful. That is, who he's a goist,' and he's nice to burn, he's like a radio, he's got to him, he's, it's caught up. "He's not right, he's just too.hit, I loved Roger. And he says it, he's got to Anyway, he was a big. He has a central, no public motivation, and he's,'s a couple of getting a shiting. Finally, I feel the words, his son, he's a teenager. Not a man, a dad he says. "I's a king he, ever he's, and he's not a battle in his years. I think it's an extrause, he's he's not he says the needs. All shews, he is the story of the 1980s, includes the game, it's a ally alike. That's what's. But it's part of the universes. It's the Army he seems to him and get ahit. Maybe it's a man that's always a point, he's eyes, a show or ahead to him. I +===== sample 2 ===== +head_tokens: ['▁to', '▁pass', '▁the', '▁Blood', '▁in', '▁itself', '▁that', '▁', 'he', '▁has', 'e', 'd', '▁that', '▁the', '▁Church', '▁of'] +tail_tokens: ['▁of', '▁the', '▁law', ',', '▁only', '▁in', '▁', 'a', '▁view', '▁with', '▁', 'a', '▁masterpiece', '.', '▁Let', ''] +to pass the Blood in itself that he hased that the Church of meaning is, he feels in a strong conscience in a confession of a casthood of the Church will tell his ideas that he thinks to see a person he tells that opposition in a hood. It is a meaning that of a Church, and in a part of the orientation that his own and son of his wife, he is with the experience of his Church, and he is entirely utter. It is at the point of view he knows his opposition to forms his own grief, and saying, he sees instead to a particular symbol of the Church suffering and the truth and that is why the name of the Church of the truth, and that spix, this is a partial part of the Church. With the steels of the Church of the Churchs to the Church of the Church's souls, the disciple religions, all of the Church, with the Church that he is from the control of the Church, as if the youth of the law. Only the source of the Emperor, in acas, and the visited champions that has a department of a juice, and become a member of a concern so that the department's head of the lores, and surrounding the Union. Under a singlehood, the member of that law, which is influenced the priest of a habit of normal, that of sins in hisglis, at the height of the founders in a sacred structure. He is aspire that reforms the king of the meaning of the priest, in the case of which the Church understanding of the fulfillment of the nation of a worthy reading to the scope of the Church, and applied with the Church of the owner of a necessary constant, ward as a unit of the wife of the victim. By the night, that he is a popular uttered for sacrifices, and equals the true DNA of the life of the Holy Primes the control of his experience and knowledge, that, on the map of the Catholic and powers of the Pope of servant, the faith, and the law grant a heirs to the Church of the time, through their actions, and the height of the commons of the field of the Church. There is a text from that of the Churchs, a number of the channels and falls in a pre-sance in a special call of the law. Only as a result, by the guilt of a slimes, and a fine understanding of a mans tradition to the account of the Church's foundation's, and the Churchs are in avier of the general Parliament and the Order of the Constitution. Only the mind the pens of the Pope's register to the ill man. There, the priests Gods of a rage and a member of the egos. According to the origin of the evils, the feder of the Church, and the heirs with the 50s of the heavens some of the priests of providing the Church of the priest. The Church is a minister of the confessions and Christ. At the hands of the dictates of the feeling of authority to experience, hands of the knowledge: a ody and a separate sum. The man he ill in the floor of the union. This he is reading of the text reads the experience of affection of the powerounce. The criticism of the priests, he tend to detail a number of the Church of the center of the modern king. This is a astic portion of the priesthood, is a a part of a lack of a church, which it submits to the poor of a priest, with the friendships and the failure of the beings. When a er from the Church of the Church, a significant awareness of a horse of arite a crast of concern: a life of the hol a yield to the Church, he treated in a disappear, and even as a trait of law. Short dunces a destroying the experience of the reform to the members of the law. Sometimes he's a simple speech, having the saads and friendship in the Truths of the the Church of the crown organs. When the priests and the priests of pastors, who is due to battle with the hands of the Church. rupting a position in the end of the Church. Since a punishment, as 13 to represent the center of the Order of the Church, a graphic member of the law, only in a view with a masterpiece. Let +===== sample 3 ===== +head_tokens: ['▁the', '▁meaning', '▁of', '▁the', '▁divine', '▁of', '▁classes', '▁of', '▁the', '▁hands', '▁of', '▁Christ', '.', '▁The', '▁view', '▁of'] +tail_tokens: ['▁law', '.', '▁The', '▁best', '▁order', '▁to', '▁the', '▁function', ',', '▁the', '▁experience', '▁is', '▁', 'a', '▁presence', ''] +the meaning of the divine of classes of the hands of Christ. The view of the texts, in a category to the the division of the meanings of the display of understanding upon the existence of the truth and insistic account of the opening of the truth Truth of itself. The the Church's hands of the word because of the extension of communion in the portion of the gate from the collapse of the Church is understood as a to the account of the Assembly and authority, the translationss of a return to the control of the Church. At the eye of pen, it is the understanding of the understanding of theoretical and relation of the approval of the Church, and the understanding of the extension of the doors of the account of the Church of the Church, and the limited actions of the understanding of the voice in the understanding of the ceiling and the view of the provision, and value up to the costs of the answer, spite of the contribution. Some of the mind, a view of the completes of the text, of the truth of the sters of the authority, is the power of the sters of amendment. Some of the part of theration of reforms, the view of the helpfulative understanding of the means of the the understanding of the practice of the text in the perspective of the account and theprien of the presence of the situation of thence in the shadow and provision of favor. Only in the organization of the preene of this figure is a power of the authority to a believers. The whole practice is as a consciousness of the shadow the part of the files, in the battle of the meaning oforthodox the worship of the truth, shadows of the practice of the truths of the Consul. But the referred of the formation of the significant matter of the shadow of a understood and the critics of the attention. The a portion of the understanding of the meaning of the truthes, is an identical to the meaning of the dealings upon theray of the matter that is followed. With the possibility of a agrees of the truth as the experience of the truths of the word and the reaches index, although of the conviction of the Prayers of the sound of the particular victim as a meaning of the constants, a copy of the evils itself and the understanding of the the fines of a formation of the truth, of the truth in the hands, the partial of the mirrors of experience, the understanding of the matter of the means and regardless of the victim, and representing the source of the relations and the youth of the interpretation, and of his experience, and understanding the grace of a common deepen in the eyes of the hands. The word of the gates is to the hereas of the Mention of the God, and a photo of the meaning of night. The symbol from the understanding of the character of exchange of the book is the rhetorics of the statue of the immediate symbol of a portion on the panel of the location of the view is the evening of a word of understanding the necessarys of a sense of a dislike of the conditions, with the movements of light the heavens. At all the truth of a moment to the center of the contents of a pastor of a gate of the neck. Around the sources, and the defenders on the formation of the truth of understanding of the museum, and the pen truth of the sum of abes of the a noble entrance. The experience of the king of the truth remains in the neither view, he described from the text. Any the relation to fighting the end of the view of reality, the passage means of a thefile and ae of the the provision oftinehood. In the central understanding of the main understanding of the prides of the man, a permanent concentration, including a certain view of the Latin, and the force of the truth of the Greek translations and the portion of the ae does, it has to be a deep perspective. It is the horse of the failures of the the meaning of the mind and the truths and actions of a connection of mind and relevant to the understanding of the rhetoric lacks the truth of the truth. For here in the mind of worship is the rhetoric, including the hands of the awareness. Christ, the loss of the commons of the moment of the mind of life. As the practice of the Church as a transmission, as the outcome is a psychological power of height. When the request of the sense of the means of the friendship, relevant to the the covering of the a battery. The explanations the moment of the rule of a deep is imposed as a the Ghost of the evening of Nowadays, a return to the heavens in charge of the law. The best order to the function, the experience is a presence +===== sample 4 ===== +head_tokens: [',', '▁but', ',', '▁to', '▁feature', '▁', 'a', '▁portion', '▁of', '▁their', '▁inequality', '.', '▁The', '▁town', '’', 's'] +tail_tokens: ['s', '▁of', '▁the', '▁man', '▁', 'a', '▁writer', '▁of', '▁the', '▁Church', 's', ',', '▁with', '▁the', '▁self', ''] +, but, to feature a portion of their inequality. The town’s decision on this map, as viewed to the vision of the city of the Canadians. The scale of democracy it’s not only the horrors of a commonly property. It is not always, according to a motion of the state of the origins of all in the country of their own legislative. The vast vision of the text is the source of the office of the state population, in the name of the state and the state’s dozens of the significance of the department. One of the integrity of the governor’s value, it’s the evidence of a conflict, fit the lines of the state of citizen, a division that of a method of England by a documentance, the nation’s in the scope of the Bible, in its view. But it’s the understanding of mind is a psychological spite of the es of the defenders of the Constitution. Deaths in the district that’s the reflection of the text. Only the presence of the president’s connections to the center’s question. It’s this is the theGA of the presence of the nation’s visions. That’s explanation, the state’s date entirely, the author of the House of England’s excesss of the independent pleasure and laws. Just the royally to the center of the secret of the provinces. Only as the basis of the state’s construction phase captures the lack of control in the state of the Constitution in a place to the district in the state that’s a principle of the state of reforms in the presence of the state’s initial state. But the a bit of explanation as the contrasts of the heads of the entertainment, of this is not a problem. With a woman’s dream and a steep vision of a larger Lord. While objective of the classics, it’s not a patient in the coming convention. In the hand of the’s in the pres, there’s a collision of a mentes the tradition of thenian Republic. It’s a little tool that in a copy of the historians, neither of the truths of the left-casts on the top of the argument. It’s a management of this, it’s the view of a one’ history of the one’s point, and one that’s rank of the dies in the fact, it’s only a point of view, if it’s upon a rule of one’s grand critics and he learns a state of evils. s the reputation markedly to the end, he’s drawn of the citizens of the one that’s the book. In the70s, a stretch of the two truth, has been in the history of a description of a class-ad of life in the future, and a twist on the narrative’s relief of the truth. It’s corresponded to the truths of Oz, he’s masses, as the 1960s of meaning the Catholics and a possibility to the truth of style. By then, there is a spite of the workings of the universe, and there is a moment of their story. In all, the father’s a book with the falls of the core of the truth of truth as a result of the expansion of the particular story. He’s seen a piece of the full specvers of the universes on the truth of the power and his god. In the end, the lack of the Catholics in the hands of the system, who’s as a matter today, and a significant element of the uttersly. It’s the rages of the tortures a page of comparison, in a sound of Catholics and the end of the author. Some of the Bible’s quite ayed, of the Church of the mind in a different understanding of the aspects of being the truth of the Lord. It’s the misunderstandings. It’s the clear in the book that the slaves are drawn to the ens of the enquisical presence of ancient rituals in the Greeks. One of the universe of the Greeks as again, and it appears to keep the simultaneously and discipline of the reads. But the heart of the Church of the height of the Queens. It’s the idea of it’s the DNA as the eyes of the planet, and the authentic. It is decided to be in the beginning of chaos, the discipline of the power of the Christ. With a position of the fan, a judgment of the one’s of the man a writer of the Churchs, with the self +===== sample 5 ===== +head_tokens: ['▁all', '▁of', '▁the', '▁wishes', '▁of', '▁the', '▁state', '▁of', '▁affection', '▁unit', 'e', 'd', '▁the', '▁metaphor', '▁of', '▁'] +tail_tokens: ['s', '▁listen', '▁to', '▁the', '▁city', '.', '▁', 'Historically', '▁', 'he', '’', 's', '▁', 'a', '▁', ''] +all of the wishes of the state of affection united the metaphor of a village that’s the finalacy in the manner of the head of the state. It is a unit of the system and the purpose of the state. The rest of the state is classified in the conditions of the Agman’s democracy. In the date of the state is driven, and figures on the Bible’s of the places and the top of a the nation’s national society. In the contrast, the national classes in the era the Church’s the touch of choice of the Resolution itself, on the other part of the initials: the top of the punishment, and the depths of the interior and the size of the citizen’s functions of the author’ss. One of the provinces contained covered on a platform, and one of the Words of the seats. The poverty of tradition is a source’s a few more. Suddenly, this is a cross-factorant, and on the affection’s a method, in the entire state, as the expansion of all the poor and the pensions of a strategy the lack of the objective of the pictures is, adding to the role of the failures of a bit, but it is out to the ages of the law. Because of the inquiry of the debatehead, the total lack of a type of document, drassed in a state, as a reformer of law, and a crew of the legislaturemens of the powers of the priest’s works, and on the margins of the state. It’s never doubt that this is meaning, that that a sla’s covered, worthy from a state council in the courts. On the Bills of the law begins, the first elections, the boy recorded a large raid of the state a lot of the MPs. As a result, he’s not still a popular conviction of a father’s of the Weeks along the area. Now it’s easy, but in a meaning of a character’s the mouth of a state tool’s body bones, while fares in the state of one of the councils that he’s involved in a Biblical war. In the dream, a whole out of the city’s eightes the chiefs of the rebel’s homeland of the nation, the part of the nation’ssts. But that’s the capital of a planet. Unlike the neighborhood is, mourn, however, a monsterprimarily, in some of the tales, it’s a lack of anger, in the personality of the sath of the rays of the stands. Yet, he’s tell about, as a man at the top of the window. One of a fellow, some of him on a moment of the life, regardless, he’s a democracy, he sees the real emerges, highlightings the sides of the earthquakes and informing the streets. The sight of his vision is in a hy of the masses. It’s also looks to be expanded in the shadow of the Waters, he that’s the matters of the case. Not he’s to become the vision of the system’s homes, running on the figure of the army in adams of the Bostons. It’s not it’s a fruit a piece of ass nash, the mythation of the state’s view. But the she’s appeared to a labs. On a long-term, a metric he is a big top-ched shadow, and the regarded. pleasure is a bit of a car, as a part of his fellow dictt. But so, in the rest of this, it’s a nice reaction. Now to develop a quick off of the corner, it’s a loss to a thing in the’s ship, but it’s a real thing to imagine, but as a sort of an n, and that’s a rock. He’s going to have a guy who’s just in control of the underground writer’s view, he’s been a bit in a child’s the tail of the car, as by a pivotal box, a foots out the city’s car and the heads of the sport. He’s the pleasure of the movie in mind as a cigarette of the city. It’s the way that he is accused of him. Some of the cases, he’s listen to the city. Historically he’s a +===== sample 6 ===== +head_tokens: ['▁experience', '▁of', '▁the', '▁same', '▁divine', '▁of', '▁classes', '▁or', '▁', 'he', 'e', 'd', '▁by', '▁the', '▁Church', '▁of'] +tail_tokens: ['▁the', '▁function', '▁of', '▁the', '▁world', '▁of', '▁the', '▁', 've', 's', '▁of', '▁the', '▁gate', '▁in', '▁the', ''] +experience of the same divine of classes or heed by the Church of the narrative, he in the Church as a sense, and a source of a anotes to the truth of the Church of itself under the Church of the writings of repeat, and as a motion of the illries, the ongoing expression of the portion of the gates principal collapses from the account of the beings of the spirit of the Orthoggs, and then the particular acknowledges of a particular constant translation truth of the text of the constant account of the transmission, a legacy of thecans and meanings. By the law of a the account of account of the operation of power, by the admiss of the Church of the current understanding of the Constitution of the Bishop, and the understanding of the law, in the view of the spite of theative valley of the complete statutenets ster. It is apparentconstituted the convictions of the Elis of the truth of the center of the destruction of theed of the general authority of the present scope of the constants in the village of the achievement of his interest to the experience of a century. The more aspect of the guidance of understanding of the background of theacts of the ts of the eye, eventually compared to the hands of dynamic and the situation of thence of the elect, is a widespread support of the current master of the dozenes, and a removal of theclaving to the representation of the dus of the card. detail is a detail of the account of the ned upon the ned of the system of the auction model of the poor and the main account of the means of meaning in the merely a member of the subject to the situation of the property. Conion as a factor of conviction of the removing the entry of the human translations, in the collision of the diversity, along with the power of the ens of the truth, and the awaits of the understanding of the association, and the text of the derivatives of the es of the powers of the Church and state of the Museum of the tents of thetes of the bread the victim of the faithful of the connection of the forces with the powers of the village of the three territories. The the unprecedented battles, with the operation of the Catholicfall of the forced forces of the damnes of the Book of the offense. By theture of the pastors of the fact of the tender in the end of the virtue of the portion of the youths, in Church, by the battle of the church and scope of the representatives of the forces of the impulses, because of the pastor and description of the attention of the writings of truth of the object of the account of relations of the Constitution. This in the understanding of the scopes of the inner of the houses of the darks, due to the truth of those in the truths of the the forces of the cause of the forces, upon of the legacy of the number of property, and the partials of the Christs is to receiving in the crown of the hands of the night of the similarities of the moment of the truth of the conviction of the formations of the contents of the heaven. As its apparently explanation, ... in the legacy of the evidence of the explanations of the formation of the no masters of a deviation of presence is a clear relation of the explicit alphabets in Christ with the the Crown of the more truthly as a portion of the center of the present expression of the time of the text. The correct understanding of the state of the competitors of the the conditions of the explanations, as a reflection of the hands of the inens and the understanding of the explosion of bones of the regimes and the absence of the Saens of the authority. Following the designer of the preenatives of the operation, drawn to the Church of translation. On the background of servants of the inds of the partiales of the authority of authority of the truth’s count understanding of the Book of life. On the combination of the vi of the convictions of the state of the mind of thedox ritual is a matter of practice as a significant appearance of the exchange of the house of the Holy writer of text andclams to the extension of the wards of the letter of the experience of order and the meaning of the evidence, in the undict. Through the extension of the power of the myths and excludes the save of the function of the Battle realization of the abolishs of the cess, and tohers the property of the top of all the deep squares and from the hands of the battle of the awareness of truth as a part of the truth. The extension of the function of the world of the ves of the gate in the +===== sample 7 ===== +head_tokens: ['▁overlook', '▁of', '▁the', '▁original', '▁death', '▁of', '▁control', '▁of', '▁the', '▁immediate', '▁detail', ',', '▁and', '▁the', '▁questions', '▁of'] +tail_tokens: ['.', '▁Yet', '▁the', '▁function', '▁of', '▁the', '▁truth', '▁of', '▁the', '▁an', 'arch', 's', '▁in', '▁the', '▁final', ''] +overlook of the original death of control of the immediate detail, and the questions of penalties of that, in the hands of the present manner. There are part of a significant understanding of the Parts in the preparation of the understanding of Christ and the part of the forgotten in the errors of the ongoing truths in the finals. The order is a display of view, and in the principality of the original understanding of the control of view of the meaning of reading up the current Treasure of failures, with the forces of the constant consciousness of the formation. By in the account the account of the understanding of the demands of testament to the meaning of the approval of the capital of the apparent truthnotes in the teacher of the ortreated of the account of the Church of the s in the Popes of the voice in the understanding of the account as a notes of the prominents of the las. The answer, a confidence in the next account of the state of the Elield of in the letter of view, and the constant sum of the Church of the general authority of the law. The Holyity of the shadow of the realm of those in the evidence of the beings of the connection of the truth of the human size of the extraordinary details of the system. The date of the theories of the account of the Constitution of the teacher of dynamic, up to the nce of the Contention of the representatives of the sources of the master of theie of the property, and the contributions Questioning to the official account of the thetion of the Though detail, in the eyes of the account of Mormon, the immediate number of the priester of the edition of the led to the forces of the sources of the sources and the end of the morning in relation to the friendship of the constituents, and the cause of the property believers of the system contrasts the evidence of the Constitution of addressing the formation of service. addition, in control the result of the interior of understanding the remains of the provision in the Book of the pleasure of forgotten of the forces of the possibility of ounce upon the attention of the meanings of the truth and support of the Royal part of a index and the account of the marked tents to the truth of the Church of the tapes of the extension of the idea with the death of the concerned Christward, and without the understanding of the control, the electative of the extension of thefessions of the freedom. The the owner of the Church is not of the rhetoric of the attention of the authoritys and the 1920s. But on the thee of the dus of the youth of the spectrum of the crime in the size of the question of the gravitation of the forces of thekovs, because of the counteative of the representatives of the Christ of the divisions of the fault of God the relations, and the leadership of God of the system is organisationd as the evidence of parallel, and in the hands of the exchanges of the matters of the exchange to apply of the limits of the panel of the gate, which of the ease of speaking to the source of the theoretical whistles of protecting the truth to the receiving of the requiredes of the Church, from the ongoingimposed of the the concentration of the Book of the justices, with the contents of the account of the completes upon, as due to the legacy of the evidence of the truths view of the situation of the masters upon: with the constructions of the Orthodoxe in the disordinations., with the Constitution understanding of the regime of the authoritys of the truth of property of the royal of the unmonitice of the Hil correct, in the range of the Vicdox forces, in the grands of the bishops of the imposed in the request of the Crown, the understandings the means of the sums the power of the account of the Greek Christs of the actual issue, and in the absence of the officialation of the force of the wins of the the translation, on the background of the account of the receiving of the Constitution of the actuals of the authority of the realm of the ens of the Eastern enemy, as the as 21s of the viative. While the intense fras of the backgrounds of the remains, of the courts of the Church. Through the original death of thedom of the Book of the forward of the popular translation, as the extension of the feature of the referrede of the consciousness of the forces of the preward within the gate of the transmission. Note the extension of the description of the Eldersters of the entire law of the law, the hands of the gate, and thever the injustices of the extension and to the formation of the canen in the background of the Greek understandings of the hands of the conditions of the Constitution of the death in the Book of the awareness of the 13 question. Yet the function of the truth of the anarchs in the final diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/logs/noise_geometry_fixed_selfcond_ce_main_pow3.log b/LTA_openwebtext_dualt/mini_owt_logdirichlet/logs/noise_geometry_fixed_selfcond_ce_main_pow3.log new file mode 100644 index 0000000000000000000000000000000000000000..47da0df36e2fa2ab17e012f1d5d0ec35e26bad62 --- /dev/null +++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/logs/noise_geometry_fixed_selfcond_ce_main_pow3.log @@ -0,0 +1,261 @@ +[ckpt] runs/owt_t5_llmclean_qwen36_35b_articlefull_10k_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_40k_elfopt_t5embed_fixed_selfcond_ce_20260529_235541/step_002000.pt +[row] step_002000 pow3 t=0.000 c=1.00 ce=6.9830 acc=0.0409 +[row] step_002000 pow3 t=0.020 c=1.00 ce=6.9808 acc=0.0413 +[row] step_002000 pow3 t=0.040 c=1.00 ce=6.9726 acc=0.0424 +[row] step_002000 pow3 t=0.060 c=1.00 ce=6.9525 acc=0.0454 +[row] step_002000 pow3 t=0.080 c=1.00 ce=6.9219 acc=0.0494 +[row] step_002000 pow3 t=0.100 c=1.01 ce=6.8708 acc=0.0555 +[row] step_002000 pow3 t=0.120 c=1.01 ce=6.8077 acc=0.0633 +[row] step_002000 pow3 t=0.140 c=1.02 ce=6.7267 acc=0.0727 +[row] step_002000 pow3 t=0.160 c=1.03 ce=6.6298 acc=0.0836 +[row] step_002000 pow3 t=0.180 c=1.04 ce=6.5191 acc=0.0967 +[row] step_002000 pow3 t=0.200 c=1.06 ce=6.3882 acc=0.1112 +[row] step_002000 pow3 t=0.220 c=1.08 ce=6.2493 acc=0.1268 +[row] step_002000 pow3 t=0.240 c=1.10 ce=6.0975 acc=0.1438 +[row] step_002000 pow3 t=0.260 c=1.13 ce=5.9212 acc=0.1626 +[row] step_002000 pow3 t=0.280 c=1.16 ce=5.7472 acc=0.1821 +[row] step_002000 pow3 t=0.300 c=1.21 ce=5.5434 acc=0.2045 +[row] step_002000 pow3 t=0.320 c=1.25 ce=5.3278 acc=0.2278 +[row] step_002000 pow3 t=0.340 c=1.31 ce=5.0961 acc=0.2531 +[row] step_002000 pow3 t=0.360 c=1.38 ce=4.8456 acc=0.2806 +[row] step_002000 pow3 t=0.380 c=1.46 ce=4.5896 acc=0.3103 +[row] step_002000 pow3 t=0.400 c=1.56 ce=4.3180 acc=0.3416 +[row] step_002000 pow3 t=0.420 c=1.67 ce=4.0419 acc=0.3752 +[row] step_002000 pow3 t=0.440 c=1.80 ce=3.7634 acc=0.4093 +[row] step_002000 pow3 t=0.460 c=1.96 ce=3.4769 acc=0.4456 +[row] step_002000 pow3 t=0.480 c=2.15 ce=3.1880 acc=0.4835 +[row] step_002000 pow3 t=0.500 c=2.38 ce=2.9057 acc=0.5217 +[row] step_002000 pow3 t=0.520 c=2.65 ce=2.6252 acc=0.5607 +[row] step_002000 pow3 t=0.540 c=2.98 ce=2.3657 acc=0.5983 +[row] step_002000 pow3 t=0.560 c=3.38 ce=2.1193 acc=0.6350 +[row] step_002000 pow3 t=0.580 c=3.87 ce=1.8865 acc=0.6706 +[row] step_002000 pow3 t=0.600 c=4.47 ce=1.6771 acc=0.7030 +[row] step_002000 pow3 t=0.620 c=5.22 ce=1.4909 acc=0.7328 +[row] step_002000 pow3 t=0.640 c=6.15 ce=1.3287 acc=0.7591 +[row] step_002000 pow3 t=0.660 c=7.34 ce=1.1882 acc=0.7821 +[row] step_002000 pow3 t=0.680 c=8.84 ce=1.0618 acc=0.8031 +[row] step_002000 pow3 t=0.700 c=10.78 ce=0.9542 acc=0.8210 +[row] step_002000 pow3 t=0.720 c=13.29 ce=0.8599 acc=0.8375 +[row] step_002000 pow3 t=0.740 c=16.59 ce=0.7687 acc=0.8534 +[row] step_002000 pow3 t=0.760 c=20.96 ce=0.6859 acc=0.8682 +[row] step_002000 pow3 t=0.780 c=26.83 ce=0.6099 acc=0.8818 +[row] step_002000 pow3 t=0.800 c=34.78 ce=0.5381 acc=0.8949 +[row] step_002000 pow3 t=0.820 c=45.69 ce=0.4667 acc=0.9080 +[row] step_002000 pow3 t=0.840 c=60.84 ce=0.4026 acc=0.9202 +[row] step_002000 pow3 t=0.860 c=82.17 ce=0.3422 acc=0.9316 +[row] step_002000 pow3 t=0.880 c=112.57 ce=0.2860 acc=0.9426 +[row] step_002000 pow3 t=0.900 c=156.50 ce=0.2301 acc=0.9534 +[row] step_002000 pow3 t=0.920 c=220.84 ce=0.1787 acc=0.9637 +[row] step_002000 pow3 t=0.940 c=316.45 ce=0.1303 acc=0.9734 +[row] step_002000 pow3 t=0.960 c=460.60 ce=0.0851 acc=0.9826 +[row] step_002000 pow3 t=0.980 c=681.19 ce=0.0422 acc=0.9914 +[row] step_002000 pow3 t=1.000 c=1024.00 ce=0.0012 acc=0.9999 +[ckpt] runs/owt_t5_llmclean_qwen36_35b_articlefull_10k_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_40k_elfopt_t5embed_fixed_selfcond_ce_20260529_235541/step_005000.pt +[row] step_005000 pow3 t=0.000 c=1.00 ce=6.8106 acc=0.0409 +[row] step_005000 pow3 t=0.020 c=1.00 ce=6.8092 acc=0.0412 +[row] step_005000 pow3 t=0.040 c=1.00 ce=6.8022 acc=0.0425 +[row] step_005000 pow3 t=0.060 c=1.00 ce=6.7813 acc=0.0455 +[row] step_005000 pow3 t=0.080 c=1.00 ce=6.7479 acc=0.0500 +[row] step_005000 pow3 t=0.100 c=1.01 ce=6.6997 acc=0.0559 +[row] step_005000 pow3 t=0.120 c=1.01 ce=6.6279 acc=0.0642 +[row] step_005000 pow3 t=0.140 c=1.02 ce=6.5366 acc=0.0745 +[row] step_005000 pow3 t=0.160 c=1.03 ce=6.4279 acc=0.0867 +[row] step_005000 pow3 t=0.180 c=1.04 ce=6.2985 acc=0.1006 +[row] step_005000 pow3 t=0.200 c=1.06 ce=6.1569 acc=0.1157 +[row] step_005000 pow3 t=0.220 c=1.08 ce=5.9878 acc=0.1331 +[row] step_005000 pow3 t=0.240 c=1.10 ce=5.7966 acc=0.1524 +[row] step_005000 pow3 t=0.260 c=1.13 ce=5.5957 acc=0.1726 +[row] step_005000 pow3 t=0.280 c=1.16 ce=5.3803 acc=0.1946 +[row] step_005000 pow3 t=0.300 c=1.21 ce=5.1542 acc=0.2183 +[row] step_005000 pow3 t=0.320 c=1.25 ce=4.9148 acc=0.2432 +[row] step_005000 pow3 t=0.340 c=1.31 ce=4.6641 acc=0.2706 +[row] step_005000 pow3 t=0.360 c=1.38 ce=4.3992 acc=0.3003 +[row] step_005000 pow3 t=0.380 c=1.46 ce=4.1299 acc=0.3316 +[row] step_005000 pow3 t=0.400 c=1.56 ce=3.8557 acc=0.3647 +[row] step_005000 pow3 t=0.420 c=1.67 ce=3.5640 acc=0.4020 +[row] step_005000 pow3 t=0.440 c=1.80 ce=3.2685 acc=0.4404 +[row] step_005000 pow3 t=0.460 c=1.96 ce=2.9732 acc=0.4805 +[row] step_005000 pow3 t=0.480 c=2.15 ce=2.6841 acc=0.5216 +[row] step_005000 pow3 t=0.500 c=2.38 ce=2.3960 acc=0.5636 +[row] step_005000 pow3 t=0.520 c=2.65 ce=2.1296 acc=0.6044 +[row] step_005000 pow3 t=0.540 c=2.98 ce=1.8707 acc=0.6451 +[row] step_005000 pow3 t=0.560 c=3.38 ce=1.6278 acc=0.6850 +[row] step_005000 pow3 t=0.580 c=3.87 ce=1.4151 acc=0.7206 +[row] step_005000 pow3 t=0.600 c=4.47 ce=1.2261 acc=0.7534 +[row] step_005000 pow3 t=0.620 c=5.22 ce=1.0634 acc=0.7826 +[row] step_005000 pow3 t=0.640 c=6.15 ce=0.9214 acc=0.8086 +[row] step_005000 pow3 t=0.660 c=7.34 ce=0.8065 acc=0.8301 +[row] step_005000 pow3 t=0.680 c=8.84 ce=0.7069 acc=0.8492 +[row] step_005000 pow3 t=0.700 c=10.78 ce=0.6216 acc=0.8656 +[row] step_005000 pow3 t=0.720 c=13.29 ce=0.5485 acc=0.8800 +[row] step_005000 pow3 t=0.740 c=16.59 ce=0.4830 acc=0.8935 +[row] step_005000 pow3 t=0.760 c=20.96 ce=0.4241 acc=0.9056 +[row] step_005000 pow3 t=0.780 c=26.83 ce=0.3679 acc=0.9171 +[row] step_005000 pow3 t=0.800 c=34.78 ce=0.3196 acc=0.9274 +[row] step_005000 pow3 t=0.820 c=45.69 ce=0.2727 acc=0.9377 +[row] step_005000 pow3 t=0.840 c=60.84 ce=0.2275 acc=0.9472 +[row] step_005000 pow3 t=0.860 c=82.17 ce=0.1908 acc=0.9554 +[row] step_005000 pow3 t=0.880 c=112.57 ce=0.1554 acc=0.9634 +[row] step_005000 pow3 t=0.900 c=156.50 ce=0.1220 acc=0.9709 +[row] step_005000 pow3 t=0.920 c=220.84 ce=0.0924 acc=0.9778 +[row] step_005000 pow3 t=0.940 c=316.45 ce=0.0662 acc=0.9840 +[row] step_005000 pow3 t=0.960 c=460.60 ce=0.0421 acc=0.9898 +[row] step_005000 pow3 t=0.980 c=681.19 ce=0.0201 acc=0.9951 +[row] step_005000 pow3 t=1.000 c=1024.00 ce=0.0005 acc=1.0000 +[ckpt] runs/owt_t5_llmclean_qwen36_35b_articlefull_10k_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_40k_elfopt_t5embed_fixed_selfcond_ce_20260529_235541/step_010000.pt +[row] step_010000 pow3 t=0.000 c=1.00 ce=6.7841 acc=0.0420 +[row] step_010000 pow3 t=0.020 c=1.00 ce=6.7828 acc=0.0423 +[row] step_010000 pow3 t=0.040 c=1.00 ce=6.7756 acc=0.0435 +[row] step_010000 pow3 t=0.060 c=1.00 ce=6.7527 acc=0.0463 +[row] step_010000 pow3 t=0.080 c=1.00 ce=6.7135 acc=0.0508 +[row] step_010000 pow3 t=0.100 c=1.01 ce=6.6565 acc=0.0565 +[row] step_010000 pow3 t=0.120 c=1.01 ce=6.5759 acc=0.0650 +[row] step_010000 pow3 t=0.140 c=1.02 ce=6.4787 acc=0.0745 +[row] step_010000 pow3 t=0.160 c=1.03 ce=6.3468 acc=0.0870 +[row] step_010000 pow3 t=0.180 c=1.04 ce=6.2007 acc=0.1007 +[row] step_010000 pow3 t=0.200 c=1.06 ce=6.0199 acc=0.1174 +[row] step_010000 pow3 t=0.220 c=1.08 ce=5.8302 acc=0.1353 +[row] step_010000 pow3 t=0.240 c=1.10 ce=5.6117 acc=0.1552 +[row] step_010000 pow3 t=0.260 c=1.13 ce=5.3787 acc=0.1773 +[row] step_010000 pow3 t=0.280 c=1.16 ce=5.1484 acc=0.1998 +[row] step_010000 pow3 t=0.300 c=1.21 ce=4.8939 acc=0.2251 +[row] step_010000 pow3 t=0.320 c=1.25 ce=4.6310 acc=0.2525 +[row] step_010000 pow3 t=0.340 c=1.31 ce=4.3488 acc=0.2830 +[row] step_010000 pow3 t=0.360 c=1.38 ce=4.0653 acc=0.3156 +[row] step_010000 pow3 t=0.380 c=1.46 ce=3.7645 acc=0.3510 +[row] step_010000 pow3 t=0.400 c=1.56 ce=3.4676 acc=0.3891 +[row] step_010000 pow3 t=0.420 c=1.67 ce=3.1695 acc=0.4283 +[row] step_010000 pow3 t=0.440 c=1.80 ce=2.8589 acc=0.4718 +[row] step_010000 pow3 t=0.460 c=1.96 ce=2.5643 acc=0.5148 +[row] step_010000 pow3 t=0.480 c=2.15 ce=2.2657 acc=0.5610 +[row] step_010000 pow3 t=0.500 c=2.38 ce=1.9896 acc=0.6051 +[row] step_010000 pow3 t=0.520 c=2.65 ce=1.7254 acc=0.6489 +[row] step_010000 pow3 t=0.540 c=2.98 ce=1.4856 acc=0.6903 +[row] step_010000 pow3 t=0.560 c=3.38 ce=1.2707 acc=0.7294 +[row] step_010000 pow3 t=0.580 c=3.87 ce=1.0751 acc=0.7660 +[row] step_010000 pow3 t=0.600 c=4.47 ce=0.9078 acc=0.7983 +[row] step_010000 pow3 t=0.620 c=5.22 ce=0.7666 acc=0.8264 +[row] step_010000 pow3 t=0.640 c=6.15 ce=0.6524 acc=0.8498 +[row] step_010000 pow3 t=0.660 c=7.34 ce=0.5580 acc=0.8698 +[row] step_010000 pow3 t=0.680 c=8.84 ce=0.4788 acc=0.8866 +[row] step_010000 pow3 t=0.700 c=10.78 ce=0.4129 acc=0.9010 +[row] step_010000 pow3 t=0.720 c=13.29 ce=0.3604 acc=0.9129 +[row] step_010000 pow3 t=0.740 c=16.59 ce=0.3099 acc=0.9242 +[row] step_010000 pow3 t=0.760 c=20.96 ce=0.2657 acc=0.9343 +[row] step_010000 pow3 t=0.780 c=26.83 ce=0.2259 acc=0.9438 +[row] step_010000 pow3 t=0.800 c=34.78 ce=0.1936 acc=0.9512 +[row] step_010000 pow3 t=0.820 c=45.69 ce=0.1609 acc=0.9592 +[row] step_010000 pow3 t=0.840 c=60.84 ce=0.1330 acc=0.9659 +[row] step_010000 pow3 t=0.860 c=82.17 ce=0.1085 acc=0.9720 +[row] step_010000 pow3 t=0.880 c=112.57 ce=0.0869 acc=0.9773 +[row] step_010000 pow3 t=0.900 c=156.50 ce=0.0661 acc=0.9827 +[row] step_010000 pow3 t=0.920 c=220.84 ce=0.0489 acc=0.9871 +[row] step_010000 pow3 t=0.940 c=316.45 ce=0.0339 acc=0.9910 +[row] step_010000 pow3 t=0.960 c=460.60 ce=0.0206 acc=0.9945 +[row] step_010000 pow3 t=0.980 c=681.19 ce=0.0097 acc=0.9974 +[row] step_010000 pow3 t=1.000 c=1024.00 ce=0.0003 acc=1.0000 +[ckpt] runs/owt_t5_llmclean_qwen36_35b_articlefull_10k_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_40k_elfopt_t5embed_fixed_selfcond_ce_20260529_235541/step_020000.pt +[row] step_020000 pow3 t=0.000 c=1.00 ce=6.7813 acc=0.0420 +[row] step_020000 pow3 t=0.020 c=1.00 ce=6.7806 acc=0.0423 +[row] step_020000 pow3 t=0.040 c=1.00 ce=6.7729 acc=0.0437 +[row] step_020000 pow3 t=0.060 c=1.00 ce=6.7502 acc=0.0467 +[row] step_020000 pow3 t=0.080 c=1.00 ce=6.7082 acc=0.0513 +[row] step_020000 pow3 t=0.100 c=1.01 ce=6.6470 acc=0.0574 +[row] step_020000 pow3 t=0.120 c=1.01 ce=6.5622 acc=0.0661 +[row] step_020000 pow3 t=0.140 c=1.02 ce=6.4473 acc=0.0765 +[row] step_020000 pow3 t=0.160 c=1.03 ce=6.3089 acc=0.0881 +[row] step_020000 pow3 t=0.180 c=1.04 ce=6.1278 acc=0.1032 +[row] step_020000 pow3 t=0.200 c=1.06 ce=5.9404 acc=0.1189 +[row] step_020000 pow3 t=0.220 c=1.08 ce=5.6683 acc=0.1397 +[row] step_020000 pow3 t=0.240 c=1.10 ce=5.4254 acc=0.1607 +[row] step_020000 pow3 t=0.260 c=1.13 ce=5.1573 acc=0.1839 +[row] step_020000 pow3 t=0.280 c=1.16 ce=4.8531 acc=0.2113 +[row] step_020000 pow3 t=0.300 c=1.21 ce=4.5496 acc=0.2407 +[row] step_020000 pow3 t=0.320 c=1.25 ce=4.2452 acc=0.2716 +[row] step_020000 pow3 t=0.340 c=1.31 ce=3.8911 acc=0.3099 +[row] step_020000 pow3 t=0.360 c=1.38 ce=3.5783 acc=0.3473 +[row] step_020000 pow3 t=0.380 c=1.46 ce=3.2466 acc=0.3903 +[row] step_020000 pow3 t=0.400 c=1.56 ce=2.9136 acc=0.4359 +[row] step_020000 pow3 t=0.420 c=1.67 ce=2.6019 acc=0.4822 +[row] step_020000 pow3 t=0.440 c=1.80 ce=2.2779 acc=0.5320 +[row] step_020000 pow3 t=0.460 c=1.96 ce=1.9928 acc=0.5790 +[row] step_020000 pow3 t=0.480 c=2.15 ce=1.7190 acc=0.6269 +[row] step_020000 pow3 t=0.500 c=2.38 ce=1.4724 acc=0.6720 +[row] step_020000 pow3 t=0.520 c=2.65 ce=1.2387 acc=0.7166 +[row] step_020000 pow3 t=0.540 c=2.98 ce=1.0377 acc=0.7569 +[row] step_020000 pow3 t=0.560 c=3.38 ce=0.8661 acc=0.7923 +[row] step_020000 pow3 t=0.580 c=3.87 ce=0.7139 acc=0.8251 +[row] step_020000 pow3 t=0.600 c=4.47 ce=0.5924 acc=0.8524 +[row] step_020000 pow3 t=0.620 c=5.22 ce=0.4892 acc=0.8760 +[row] step_020000 pow3 t=0.640 c=6.15 ce=0.4044 acc=0.8957 +[row] step_020000 pow3 t=0.660 c=7.34 ce=0.3401 acc=0.9113 +[row] step_020000 pow3 t=0.680 c=8.84 ce=0.2875 acc=0.9243 +[row] step_020000 pow3 t=0.700 c=10.78 ce=0.2449 acc=0.9347 +[row] step_020000 pow3 t=0.720 c=13.29 ce=0.2071 acc=0.9443 +[row] step_020000 pow3 t=0.740 c=16.59 ce=0.1761 acc=0.9522 +[row] step_020000 pow3 t=0.760 c=20.96 ce=0.1496 acc=0.9590 +[row] step_020000 pow3 t=0.780 c=26.83 ce=0.1264 acc=0.9653 +[row] step_020000 pow3 t=0.800 c=34.78 ce=0.1043 acc=0.9712 +[row] step_020000 pow3 t=0.820 c=45.69 ce=0.0858 acc=0.9760 +[row] step_020000 pow3 t=0.840 c=60.84 ce=0.0695 acc=0.9804 +[row] step_020000 pow3 t=0.860 c=82.17 ce=0.0552 acc=0.9844 +[row] step_020000 pow3 t=0.880 c=112.57 ce=0.0433 acc=0.9877 +[row] step_020000 pow3 t=0.900 c=156.50 ce=0.0324 acc=0.9908 +[row] step_020000 pow3 t=0.920 c=220.84 ce=0.0234 acc=0.9933 +[row] step_020000 pow3 t=0.940 c=316.45 ce=0.0157 acc=0.9955 +[row] step_020000 pow3 t=0.960 c=460.60 ce=0.0095 acc=0.9973 +[row] step_020000 pow3 t=0.980 c=681.19 ce=0.0044 acc=0.9987 +[row] step_020000 pow3 t=1.000 c=1024.00 ce=0.0002 acc=1.0000 +[ckpt] runs/owt_t5_llmclean_qwen36_35b_articlefull_10k_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_40k_elfopt_t5embed_fixed_selfcond_ce_20260529_235541/step_040000.pt +[row] step_040000 pow3 t=0.000 c=1.00 ce=6.7724 acc=0.0417 +[row] step_040000 pow3 t=0.020 c=1.00 ce=6.7711 acc=0.0419 +[row] step_040000 pow3 t=0.040 c=1.00 ce=6.7627 acc=0.0432 +[row] step_040000 pow3 t=0.060 c=1.00 ce=6.7389 acc=0.0464 +[row] step_040000 pow3 t=0.080 c=1.00 ce=6.7019 acc=0.0507 +[row] step_040000 pow3 t=0.100 c=1.01 ce=6.6330 acc=0.0572 +[row] step_040000 pow3 t=0.120 c=1.01 ce=6.5428 acc=0.0654 +[row] step_040000 pow3 t=0.140 c=1.02 ce=6.4288 acc=0.0760 +[row] step_040000 pow3 t=0.160 c=1.03 ce=6.2732 acc=0.0888 +[row] step_040000 pow3 t=0.180 c=1.04 ce=6.0878 acc=0.1034 +[row] step_040000 pow3 t=0.200 c=1.06 ce=5.8291 acc=0.1224 +[row] step_040000 pow3 t=0.220 c=1.08 ce=5.6040 acc=0.1413 +[row] step_040000 pow3 t=0.240 c=1.10 ce=5.2579 acc=0.1667 +[row] step_040000 pow3 t=0.260 c=1.13 ce=4.9081 acc=0.1966 +[row] step_040000 pow3 t=0.280 c=1.16 ce=4.5561 acc=0.2279 +[row] step_040000 pow3 t=0.300 c=1.21 ce=4.2154 acc=0.2628 +[row] step_040000 pow3 t=0.320 c=1.25 ce=3.7887 acc=0.3072 +[row] step_040000 pow3 t=0.340 c=1.31 ce=3.4414 acc=0.3486 +[row] step_040000 pow3 t=0.360 c=1.38 ce=3.0283 acc=0.4014 +[row] step_040000 pow3 t=0.380 c=1.46 ce=2.6739 acc=0.4517 +[row] step_040000 pow3 t=0.400 c=1.56 ce=2.2996 acc=0.5093 +[row] step_040000 pow3 t=0.420 c=1.67 ce=1.9696 acc=0.5642 +[row] step_040000 pow3 t=0.440 c=1.80 ce=1.6642 acc=0.6187 +[row] step_040000 pow3 t=0.460 c=1.96 ce=1.4070 acc=0.6681 +[row] step_040000 pow3 t=0.480 c=2.15 ce=1.1724 acc=0.7152 +[row] step_040000 pow3 t=0.500 c=2.38 ce=0.9526 acc=0.7613 +[row] step_040000 pow3 t=0.520 c=2.65 ce=0.7835 acc=0.7992 +[row] step_040000 pow3 t=0.540 c=2.98 ce=0.6332 acc=0.8340 +[row] step_040000 pow3 t=0.560 c=3.38 ce=0.5082 acc=0.8642 +[row] step_040000 pow3 t=0.580 c=3.87 ce=0.4092 acc=0.8885 +[row] step_040000 pow3 t=0.600 c=4.47 ce=0.3269 acc=0.9096 +[row] step_040000 pow3 t=0.620 c=5.22 ce=0.2643 acc=0.9259 +[row] step_040000 pow3 t=0.640 c=6.15 ce=0.2158 acc=0.9389 +[row] step_040000 pow3 t=0.660 c=7.34 ce=0.1787 acc=0.9490 +[row] step_040000 pow3 t=0.680 c=8.84 ce=0.1470 acc=0.9575 +[row] step_040000 pow3 t=0.700 c=10.78 ce=0.1231 acc=0.9643 +[row] step_040000 pow3 t=0.720 c=13.29 ce=0.1030 acc=0.9699 +[row] step_040000 pow3 t=0.740 c=16.59 ce=0.0865 acc=0.9747 +[row] step_040000 pow3 t=0.760 c=20.96 ce=0.0721 acc=0.9788 +[row] step_040000 pow3 t=0.780 c=26.83 ce=0.0596 acc=0.9825 +[row] step_040000 pow3 t=0.800 c=34.78 ce=0.0499 acc=0.9853 +[row] step_040000 pow3 t=0.820 c=45.69 ce=0.0401 acc=0.9880 +[row] step_040000 pow3 t=0.840 c=60.84 ce=0.0317 acc=0.9906 +[row] step_040000 pow3 t=0.860 c=82.17 ce=0.0253 acc=0.9924 +[row] step_040000 pow3 t=0.880 c=112.57 ce=0.0192 acc=0.9943 +[row] step_040000 pow3 t=0.900 c=156.50 ce=0.0149 acc=0.9956 +[row] step_040000 pow3 t=0.920 c=220.84 ce=0.0104 acc=0.9969 +[row] step_040000 pow3 t=0.940 c=316.45 ce=0.0070 acc=0.9979 +[row] step_040000 pow3 t=0.960 c=460.60 ce=0.0043 acc=0.9987 +[row] step_040000 pow3 t=0.980 c=681.19 ce=0.0020 acc=0.9994 +[row] step_040000 pow3 t=1.000 c=1024.00 ce=0.0002 acc=1.0000 +wrote ../docs/lta_samples/metrics_20260531/noise_geometry_fixed_selfcond_ce_main/pow3/noise_geometry.tsv